python: bind PixelHistogramCUDA

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).

include/aare/hist/PixelHistogramCUDA.hpp

@@ -62,6 +62,7 @@ class PixelHistogramCUDA {
     // The order of the declarations below is important. It conditions the order
     // of their destruction.
     HostRegistration m_registered_input{};
+    HostRegistration m_registered_output{};
 
     // Device buffers for frames and histogram
     DeviceBuffer<DeviceStorageType> d_histogram;
@@ -155,7 +156,20 @@ class PixelHistogramCUDA {
         m_registered_input.unregister();
     }
 
-    void submit_batch(NDView<InputType, 3> frames) {
+    // Pins the caller's output histogram, so values(out) into it copies at
+    // DMA speed. The D2H copy runs on the compute stream.
+    void register_output_buffer(void *ptr, size_t bytes) {
+        CUDA_CHECK(cudaStreamSynchronize(compute_stream.get()));
+        m_registered_output.unregister();
+        m_registered_output = HostRegistration(ptr, bytes);
+    }
+
+    void unregister_output_buffer() {
+        CUDA_CHECK(cudaStreamSynchronize(compute_stream.get()));
+        m_registered_output.unregister();
+    }
+
+    void submit_batch(NDView<const InputType, 3> frames) {
         auto frame_shape = frames.shape();
         if (static_cast<size_t>(frame_shape[0]) > m_batch_size ||
             static_cast<size_t>(frame_shape[1]) != nrows ||
@@ -205,11 +219,11 @@ class PixelHistogramCUDA {
             cudaEventRecord(kernel_done[slot].get(), compute_stream.get()));
     }
 
-    void fill_histogram(NDView<InputType, 3> frames) {
+    void fill_histogram(NDView<const InputType, 3> frames) {
         auto N_frames = static_cast<size_t>(frames.shape(0));
         for (size_t i = 0; i < N_frames; i += m_batch_size) {
             size_t batch_end = std::min(i + m_batch_size, N_frames);
-            NDView<InputType, 3> batch = frames.sub_view(i, batch_end);
+            NDView<const InputType, 3> batch = frames.sub_view(i, batch_end);
             submit_batch(batch);
         }
     }

python/CMakeLists.txt

@@ -76,6 +76,7 @@ set(PYTHON_FILES
     aare/Cluster.py
     aare/FastPedestal.py
     aare/Pedestal.py
+    aare/PixelHistogramCUDA.py
     aare/calibration.py
     aare/factory.py
     aare/experimental.py

python/aare/PixelHistogramCUDA.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: MPL-2.0
+import numpy as np
+
+from . import _aare
+from .ClusterFinder import _cuda_available
+
+# numpy dtype -> class name suffix, for the frames and for the counts.
+_INPUT_SUFFIX = {
+    np.dtype(np.uint16): "u16",
+    np.dtype(np.float32): "f",
+    np.dtype(np.float64): "d",
+}
+_STORAGE_SUFFIX = {
+    np.dtype(np.uint8): "u8",
+    np.dtype(np.uint16): "u16",
+    np.dtype(np.uint32): "u32",
+    np.dtype(np.uint64): "u64",
+    np.dtype(np.float32): "f",
+    np.dtype(np.float64): "d",
+}
+
+
+def _suffix(table, dtype, what):
+    try:
+        return table[np.dtype(dtype)]
+    except (KeyError, TypeError):
+        supported = ", ".join(str(d) for d in table)
+        raise ValueError(
+            f"Unsupported {what}: {dtype}. Supported: {supported}."
+        ) from None
+
+
+def PixelHistogramCUDA(*, rows, cols, n_bins, xmin, xmax, dtype,
+                       storage_dtype=np.uint32, batch_size=64):
+    """
+    Per-pixel histograms filled on the GPU.
+
+    Parameters
+    ----------
+    rows, cols : int
+        Frame shape.
+    n_bins : int
+        Number of bins per pixel.
+    xmin, xmax : float
+        Histogram range, [xmin, xmax). Values outside it, and NaN, are
+        dropped.
+    dtype : numpy dtype
+        dtype of the frames you will pass: np.uint16, np.float32 or
+        np.float64. Frames are copied to the GPU as they are and are never
+        converted, so fill_histogram() raises TypeError for any other dtype.
+        The copy usually limits the speed: uint16 frames move half the bytes
+        of float32 and a quarter of float64. Use float32 rather than float64
+        for processed frames when its precision is enough.
+    storage_dtype : numpy dtype, optional
+        dtype of the counts returned by values(): np.uint8, np.uint16,
+        np.uint32, np.uint64, np.float32 or np.float64. Integer counts are
+        capped at the type's maximum. Default np.uint32.
+    batch_size : int, optional
+        Frames copied and processed per step. Two batches are held on the
+        GPU, 2 * batch_size * rows * cols * itemsize bytes. Default 64.
+
+    Example
+    -------
+    .. code-block:: python
+
+        import numpy as np
+        from aare import PixelHistogramCUDA
+
+        frames = np.fromfile(path, dtype=np.uint16).reshape(-1, 400, 400)
+        hist = PixelHistogramCUDA(rows=400, cols=400, n_bins=1000,
+                                  xmin=-0.5, xmax=999.5, dtype=frames.dtype)
+        hist.register_input_buffer(frames)  # optional: faster copies
+        hist.fill_histogram(frames)
+        counts = hist.values()              # (400, 400, 1000)
+    """
+    if not _cuda_available():
+        raise RuntimeError(
+            "PixelHistogramCUDA is not available in this build of aare. "
+            "Rebuild with -DAARE_CUDA=ON (and -DAARE_PYTHON_BINDINGS=ON)."
+        )
+    name = (f"PixelHistogramCUDA_{_suffix(_INPUT_SUFFIX, dtype, 'dtype')}_"
+            f"{_suffix(_STORAGE_SUFFIX, storage_dtype, 'storage_dtype')}")
+    cls = getattr(_aare, name)
+    return cls(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax,
+               batch_size=batch_size)

python/aare/__init__.py

@@ -4,16 +4,16 @@ from . import _aare
 
 # ---- CUDA module (optional) ------------------------------------------------
 # When the package was built with AARE_CUDA=ON, a sibling extension
-# _aare_cuda contains the ClusterFinderCUDA_* classes. We re-export them
-# onto _aare so user code can do `from aare import ClusterFinderCUDA_*`
-# regardless of which .so physically hosts the class. On a CPU-only build
-# the import fails silently and ClusterFinderCUDA_* classes simply aren't
-# present; the factory in ClusterFinder.py handles that case with a clear
-# RuntimeError.
+# _aare_cuda contains the ClusterFinderCUDA_* and PixelHistogramCUDA_*
+# classes. We re-export them onto _aare so user code can do
+# `from aare import ClusterFinderCUDA_*` regardless of which .so physically
+# hosts the class. On a CPU-only build the import fails silently and the
+# classes simply aren't present; the factories in ClusterFinder.py and
+# PixelHistogramCUDA.py handle that case with a clear RuntimeError.
 try:
     from . import _aare_cuda as _aare_cuda_mod
     for _name in dir(_aare_cuda_mod):
-        if _name.startswith("ClusterFinderCUDA"):
+        if _name.startswith(("ClusterFinderCUDA", "PixelHistogramCUDA")):
             setattr(_aare, _name, getattr(_aare_cuda_mod, _name))
     del _name
 except ImportError:
@@ -55,6 +55,7 @@ from .ClusterFinder import (ClusterFinder, ClusterFinderFrozen, ClusterCollector
                             ClusterFinderCUDA, ClusterFinderCUDAGraph, _cuda_available, find_cluster_views_batched_iter)
 from .ClusterVector import ClusterVector
 from .Cluster import Cluster
+from .PixelHistogramCUDA import PixelHistogramCUDA
 
 from ._aare import Gaussian, RisingScurve, FallingScurve, Pol1, Pol2, GaussianErfcPlateau, GaussianChargeSharing, GaussianChargeSharingKb
 from ._aare import fit

python/src/bind_PixelHistogramCUDA.hpp

@@ -0,0 +1,237 @@
+// SPDX-License-Identifier: MPL-2.0
+#pragma once
+#include "aare/NDArray.hpp"
+#include "aare/NDView.hpp"
+#include "aare/hist/PixelHistogramCUDA.hpp"
+#include "np_helper.hpp"
+
+#include <cstddef>
+#include <cstdint>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <string>
+#include <type_traits>
+
+namespace py = pybind11;
+
+namespace aare {
+
+namespace detail {
+
+// Frames are copied to the GPU as given, so a wrong dtype must not be
+// converted silently: the converted copy would be a temporary, not the
+// caller's (possibly pinned) buffer. Ndim and contiguity are checked when the
+// view is made.
+template <typename InputType>
+py::array_t<InputType> require_dtype(const py::array &arr) {
+    if (!py::isinstance<py::array_t<InputType>>(arr)) {
+        const auto want =
+            py::str(py::dtype::of<InputType>()).cast<std::string>();
+        const auto got = py::str(arr.dtype()).cast<std::string>();
+        throw py::type_error(
+            "PixelHistogramCUDA was built for " + want + " frames, got " + got +
+            ". Build it with dtype=" + got +
+            ", or convert the frames with .astype(np." + want + ").");
+    }
+    return py::reinterpret_borrow<py::array_t<InputType>>(arr);
+}
+
+// An output histogram must be the storage dtype, shape (rows, cols, n_bins),
+// C-contiguous and writeable: values(out) copies straight into it.
+template <typename StorageType>
+py::array_t<StorageType> require_output(const py::array &arr, size_t rows,
+                                        size_t cols, size_t n_bins) {
+    const auto want = py::str(py::dtype::of<StorageType>()).cast<std::string>();
+    if (!py::isinstance<py::array_t<StorageType>>(arr))
+        throw py::type_error("out must have dtype " + want + ", got " +
+                             py::str(arr.dtype()).cast<std::string>() + ".");
+    if (arr.ndim() != 3 || static_cast<size_t>(arr.shape(0)) != rows ||
+        static_cast<size_t>(arr.shape(1)) != cols ||
+        static_cast<size_t>(arr.shape(2)) != n_bins)
+        throw py::value_error("out must have shape (" + std::to_string(rows) +
+                              ", " + std::to_string(cols) + ", " +
+                              std::to_string(n_bins) + ").");
+    if (!(arr.flags() & py::array::c_style))
+        throw py::value_error("out must be C-contiguous.");
+    if (!arr.writeable())
+        throw py::value_error("out must be writeable.");
+    return py::reinterpret_borrow<py::array_t<StorageType>>(arr);
+}
+
+} // namespace detail
+
+template <typename StorageType, typename InputType>
+void define_PixelHistogramCUDA(py::module &m, const std::string &class_name) {
+    // 64-bit device counters for 8-byte storage, so counts can't wrap at
+    // 2^32 before the cap in values().
+    using DeviceStorageType =
+        std::conditional_t<sizeof(StorageType) == 8, uint64_t, uint32_t>;
+    using Hist =
+        cuda::PixelHistogramCUDA<StorageType, DeviceStorageType, InputType>;
+    using AxisType = typename Hist::AxisType;
+
+    py::class_<Hist>(m, class_name.c_str(),
+                     R"(Per-pixel histograms filled on the GPU.
+
+             Frames are copied to the device with the dtype they have; the
+             histogram axis is floating point. Build instances with the
+             aare.PixelHistogramCUDA() factory rather than by class name.)")
+        .def(py::init<size_t, size_t, size_t, AxisType, AxisType, size_t>(),
+             py::kw_only(), py::arg("rows"), py::arg("cols"), py::arg("n_bins"),
+             py::arg("xmin"), py::arg("xmax"), py::arg("batch_size"))
+
+        .def(
+            "register_input_buffer",
+            [](Hist &self, const py::array &frames) {
+                auto arr = detail::require_dtype<InputType>(frames);
+                if (!(arr.flags() & py::array::c_style))
+                    throw py::value_error("Array is not C-contiguous. Use "
+                                          "np.ascontiguousarray(arr)");
+                self.register_input_buffer(
+                    const_cast<void *>(static_cast<const void *>(arr.data())),
+                    static_cast<size_t>(arr.nbytes()));
+            },
+            py::arg("frames"), py::keep_alive<1, 2>(),
+            R"(Pin `frames` so that fill_histogram() copies it at full DMA
+             bandwidth and overlaps the copies with the kernels.
+
+             Pass the whole dataset once; slices of it passed to
+             fill_histogram() benefit too. The array is kept alive for as
+             long as the histogram exists. Registering another array releases
+             the previous one.)")
+
+        .def("unregister_input_buffer", &Hist::unregister_input_buffer,
+             "Release the buffer pinned by register_input_buffer().")
+
+        .def(
+            "fill_histogram",
+            [](Hist &self, const py::array &frames) {
+                auto arr = detail::require_dtype<InputType>(frames);
+                auto view = make_const_view_3d(arr);
+                py::gil_scoped_release release;
+                self.fill_histogram(view);
+            },
+            py::arg("frames"),
+            R"(Queue frames of shape (n_frames, rows, cols) for histogramming,
+             in batches of batch_size.
+
+             Returns once the work is queued, not when it is done. If the
+             frames lie in a registered buffer, the copies run in the
+             background: don't modify them until flush() or values() has
+             returned.)")
+
+        .def("flush", &Hist::flush, py::call_guard<py::gil_scoped_release>(),
+             "Block until every queued frame has been histogrammed.")
+
+        .def(
+            "register_output_buffer",
+            [](Hist &self, const py::array &out) {
+                auto arr = detail::require_output<StorageType>(
+                    out, self.rows(), self.cols(), self.n_bins());
+                self.register_output_buffer(arr.mutable_data(),
+                                            static_cast<size_t>(arr.nbytes()));
+            },
+            py::arg("out"), py::keep_alive<1, 2>(),
+            R"(Pin `out` so that values(out=out) copies into it at full DMA
+             bandwidth. `out` must have the storage dtype and shape
+             (rows, cols, n_bins). It is kept alive for as long as the
+             histogram exists. Registering another array releases the
+             previous one.)")
+
+        .def("unregister_output_buffer", &Hist::unregister_output_buffer,
+             "Release the buffer pinned by register_output_buffer().")
+
+        .def(
+            "values",
+            [](const Hist &self, const py::object &out) -> py::object {
+                if (out.is_none()) {
+                    NDArray<StorageType, 3> *ptr = nullptr;
+                    {
+                        py::gil_scoped_release release;
+                        ptr = new NDArray<StorageType, 3>(self.values());
+                    }
+                    return return_image_data(ptr);
+                }
+                if (!py::isinstance<py::array>(out))
+                    throw py::type_error("out must be a numpy array.");
+                auto arr = detail::require_output<StorageType>(
+                    py::reinterpret_borrow<py::array>(out), self.rows(),
+                    self.cols(), self.n_bins());
+                auto view = make_view_3d(arr);
+                {
+                    py::gil_scoped_release release;
+                    self.values(view);
+                }
+                return out;
+            },
+            py::arg("out") = py::none(),
+            R"(Histogram of every queued frame, shape (rows, cols, n_bins).
+             Counts are capped at the storage dtype's maximum. Waits for the
+             queued work; filling can continue afterwards.
+
+             With `out`, the histogram is written into that array, which is
+             returned. Reusing one array, pinned with
+             register_output_buffer(), avoids allocating and copying through
+             pageable memory on every call.)")
+
+        .def(
+            "values_pixel",
+            [](const Hist &self, size_t row, size_t col) {
+                NDArray<StorageType, 1> *ptr = nullptr;
+                {
+                    py::gil_scoped_release release;
+                    ptr = new NDArray<StorageType, 1>(
+                        self.values_pixel(row, col));
+                }
+                return return_image_data(ptr);
+            },
+            py::arg("row"), py::arg("col"),
+            R"(Histogram of one pixel, shape (n_bins,). For many pixels, call
+             values() once instead. Raises IndexError outside the frame.)")
+
+        .def(
+            "bin_edges",
+            [](const Hist &self) {
+                return return_image_data(
+                    new NDArray<AxisType, 1>(self.bin_edges()));
+            },
+            "Bin edges, n_bins + 1 values from xmin to xmax.")
+
+        .def(
+            "bin_centers",
+            [](const Hist &self) {
+                return return_image_data(
+                    new NDArray<AxisType, 1>(self.bin_centers()));
+            },
+            "Bin centres, n_bins values.")
+
+        .def_property_readonly("rows", &Hist::rows)
+        .def_property_readonly("cols", &Hist::cols)
+        .def_property_readonly("n_bins", &Hist::n_bins)
+        .def_property_readonly("batch_size", &Hist::batch_size)
+        .def_property_readonly(
+            "dtype", [](const Hist &) { return py::dtype::of<InputType>(); },
+            "dtype of the frames this histogram takes.");
+}
+
+// Class names are PixelHistogramCUDA_<input>_<storage>, e.g.
+// PixelHistogramCUDA_u16_u32 for uint16 frames and uint32 counts.
+template <typename InputType>
+void define_PixelHistogramCUDA_storages(py::module &m,
+                                        const std::string &input) {
+    const std::string base = "PixelHistogramCUDA_" + input + "_";
+    define_PixelHistogramCUDA<uint8_t, InputType>(m, base + "u8");
+    define_PixelHistogramCUDA<uint16_t, InputType>(m, base + "u16");
+    define_PixelHistogramCUDA<uint32_t, InputType>(m, base + "u32");
+    define_PixelHistogramCUDA<uint64_t, InputType>(m, base + "u64");
+    define_PixelHistogramCUDA<float, InputType>(m, base + "f");
+    define_PixelHistogramCUDA<double, InputType>(m, base + "d");
+}
+
+inline void define_PixelHistogramCUDA_bindings(py::module &m) {
+    define_PixelHistogramCUDA_storages<uint16_t>(m, "u16");
+    define_PixelHistogramCUDA_storages<float>(m, "f");
+    define_PixelHistogramCUDA_storages<double>(m, "d");
+}
+
+} // namespace aare

python/src/cuda_bindings.cu

@@ -4,11 +4,13 @@
 // 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>
 
@@ -83,6 +85,9 @@ 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);
+
+    // Pixel histograms: PixelHistogramCUDA_<input>_<storage>
+    aare::define_PixelHistogramCUDA_bindings(m);
 }
 
 #undef DEFINE_CUDA_CLUSTER_TYPES

python/tests/PixelHistogramCUDA.ipynb

@@ -0,0 +1,224 @@
+{
+    "cells": [
+        {
+            "cell_type": "markdown",
+            "id": "9b1ef299",
+            "metadata": {},
+            "source": [
+                "# PixelHistogram — CPU vs CUDA\n",
+                "\n",
+                "Per-pixel histograms of a data file, filled by `PixelHistogram` on the CPU and\n",
+                "by `PixelHistogramCUDA` on the GPU. Times both, then checks that the results are\n",
+                "identical."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": 1,
+            "id": "6c8b158f",
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import sys; sys.path.append('/home/ferjao_k/aare/build')\n",
+                "\n",
+                "import time\n",
+                "from pathlib import Path\n",
+                "\n",
+                "import matplotlib.pyplot as plt\n",
+                "import numpy as np\n",
+                "\n",
+                "from aare import File, PixelHistogram_u32, PixelHistogramCUDA"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": 2,
+            "id": "518fd5c4",
+            "metadata": {},
+            "outputs": [
+                {
+                    "name": "stdout",
+                    "output_type": "stream",
+                    "text": [
+                        "(100000, 400, 400) uint16 32.00 GB\n"
+                    ]
+                }
+            ],
+            "source": [
+                "base = Path('/mnt/sls_det_storage/moench_data/2603_MaxIVBeamtime/2026032408/process/xrf/')\n",
+                "f = File(base / 'Cu_factor_10_data_master_0.json')\n",
+                "frames = f.read_n(100_000)          # (n_frames, rows, cols), uint16\n",
+                "n_frames, rows, cols = frames.shape\n",
+                "\n",
+                "n_bins, xmin, xmax = 1000, 3500, 7500   # 4 ADU per bin\n",
+                "n_threads = 16                          # CPU threads\n",
+                "print(frames.shape, frames.dtype, f'{frames.nbytes / 1e9:.2f} GB')"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "id": "7c40949d",
+            "metadata": {},
+            "source": [
+                "## CPU"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": 3,
+            "id": "a418bbd4",
+            "metadata": {},
+            "outputs": [
+                {
+                    "name": "stdout",
+                    "output_type": "stream",
+                    "text": [
+                        "CPU: 21.75 s, 4,598 frames/s\n"
+                    ]
+                }
+            ],
+            "source": [
+                "cpu = PixelHistogram_u32(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax,\n",
+                "                         n_threads=n_threads)\n",
+                "\n",
+                "t0 = time.perf_counter()\n",
+                "for frame in frames:\n",
+                "    cpu.fill_async(frame.astype(np.float64))   # the CPU binding takes float64\n",
+                "values_cpu = cpu.values()\n",
+                "t_cpu = time.perf_counter() - t0\n",
+                "print(f'CPU: {t_cpu:.2f} s, {n_frames / t_cpu:,.0f} frames/s')"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "id": "0b95f098",
+            "metadata": {},
+            "source": [
+                "## GPU"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": 4,
+            "id": "a60415d8",
+            "metadata": {},
+            "outputs": [
+                {
+                    "name": "stdout",
+                    "output_type": "stream",
+                    "text": [
+                        "GPU fill:   1.512 s, 66,119 frames/s, 21.2 GB/s\n",
+                        "GPU values: 0.032 s\n",
+                        "GPU total:  1.545 s, speed-up 14.1x\n"
+                    ]
+                }
+            ],
+            "source": [
+                "gpu = PixelHistogramCUDA(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax,\n",
+                "                         dtype=frames.dtype, storage_dtype=np.uint32, batch_size=128)\n",
+                "gpu.register_input_buffer(frames)   # pin once, before the timed loop\n",
+                "values_gpu = np.empty((rows, cols, n_bins), dtype=np.uint32)\n",
+                "gpu.register_output_buffer(values_gpu)\n",
+                "\n",
+                "t0 = time.perf_counter()\n",
+                "gpu.fill_histogram(frames)\n",
+                "gpu.flush()                         # fill_histogram only queues the work\n",
+                "t1 = time.perf_counter()\n",
+                "gpu.values(out=values_gpu)          # transpose + copy back to the host\n",
+                "t2 = time.perf_counter()\n",
+                "\n",
+                "t_gpu = t2 - t0\n",
+                "print(f'GPU fill:   {t1 - t0:.3f} s, {n_frames / (t1 - t0):,.0f} frames/s, '\n",
+                "      f'{frames.nbytes / (t1 - t0) / 1e9:.1f} GB/s')\n",
+                "print(f'GPU values: {t2 - t1:.3f} s')\n",
+                "print(f'GPU total:  {t_gpu:.3f} s, speed-up {t_cpu / t_gpu:.1f}x')"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "id": "43fff4d8",
+            "metadata": {},
+            "source": [
+                "## Compare"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": 5,
+            "id": "014cfc32",
+            "metadata": {},
+            "outputs": [
+                {
+                    "name": "stdout",
+                    "output_type": "stream",
+                    "text": [
+                        "identical: True\n",
+                        "counts:    15997375206 15997375206\n"
+                    ]
+                }
+            ],
+            "source": [
+                "print('identical:', np.array_equal(values_cpu, values_gpu))\n",
+                "print('counts:   ', values_cpu.sum(), values_gpu.sum())"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": 6,
+            "id": "a0571bd1",
+            "metadata": {},
+            "outputs":
+                [
+                    {
+                        "data": {
+                            "image/png": "iVBORw0KGgoAAAANSUhEUgAAA/YAAAGHCAYAAAAa+ckZAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAPYQAAD2EBqD+naQAArmFJREFUeJzs3Qd8E+UbB/BfV1pm2UuGIHuq7CV7yRBRQUQEBRTZUkSGLaVlCMhQhgIqqEz/KriRIbL3kKUIsqHsXWiTpvl/njfNNemgK23W7/v5HJe7XNPLhebuufd5n9fLZDKZQEREREREREQuydvRO0BERERERERE6cfAnoiIiIiIiMiFMbAnIiIiIiIicmEM7ImIiIiIiIhcGAN7IiIiIiIiIhfGwJ6IiIiIiIjIhTGwJyIiIiIiInJhDOyJiIiIiIiIXBgDeyIiIiIiIiIXxsCeyIk1bdpUTZll8eLF8PLywpkzZ1K1fVhYGCpXrozY2Fi1HBERgffffx/169dHgQIFkDt3btSsWRMLFiyA0WhM9PP379/HsGHDUKxYMQQEBODJJ5/EihUrkvxd+/fvR8uWLZEzZ07kyZMHXbp0walTp9L9Xr///nt0794dZcuWRbZs2fD444+jR48eOHHiRJLbr1+/Xr2v7Nmzq/fWu3dvXL16NdF2BoMB48ePV6/n7++PihUrYvbs2Ym269mzJzp37pzu/SciIrIXOe/L+V+uA9J7TeCK1zFE7szX0TtARMmbN2+e0xyeS5cuYerUqeok6u1tvie4b98+fPXVV3jttdcQHBwMPz8//Pbbb3j77bexc+dOfPHFFzavIcH5nj178MEHH6B8+fJYtmyZCrblRsErr7yibffPP/+oCwEJ/L/55htERUUhJCQEjRs3xsGDB1GwYME07/+UKVNQpEgRjB07FmXKlMH58+cxadIkPP3002pfq1Spom27adMmtGvXDu3bt8cPP/ygAvr33nsPLVq0wN69e1UAbzFgwAB8/fXXCA8PR+3atfH7779j6NChuHfvHsaMGaNtFxoaqoL+P/74A82bN0/z/hMREbkaZ7qOIXJ7JiLyWIsWLTLJ18Dp06dT3HbkyJGmxx57zGQ0GrV1N2/eNOn1+kTbDhw4UL3uuXPntHW//PKLWrds2TKbbVu1amUqVqyYKSYmRlv30ksvmQoUKGC6c+eOtu7MmTMmPz8/tR/pceXKlUTrLl68qF6zT58+Nutr165tqly5sslgMGjrtm3bpvZ/3rx52rojR46YvLy8TJMmTbL5+X79+pmyZctmunHjhs36Dh06qPdLRETkSHLel3OaXAek55rAWbjiPhNlFqbiE2UxabmVtLEDBw6oFmxJXw8MDMSrr76Ka9euPTKFTVq6pbX8p59+stlO0sQlZfzw4cM2qeTSwiyvL881bNgQGzZsSNc+6/V6fP7556pV3dJaL/Lmzata6ROqU6eOml+4cEFbt2rVKpVW/9JLL9ls+/rrr6tsgF27dqnlmJgY/Pzzz3jhhRfUvluUKlUKzZo1U6+THoUKFUq0TroEFC9eXLXeW1y8eFFlFUjqvK9vfFJTgwYNVJaB9e9fvXq13BxV7yHhe3r48CHWrFljs15eUz6X//77L13vgYiI6FFOnjypzkHlypVT5/7HHnsMHTt2tLk+cPfrGNmvDh06qPO+ZNjJuV4y8KyvSYjcEQN7Igd5/vnnVX/vb7/9Vp0kJUhs06aN6rOdHEkHlxTxXr164ezZs2rdokWL8OWXX6p+3dWqVVPrlixZgtatW6uToTwn6ez58uVTr5+e4F6C7hs3bqjAOjUk3VyCYgmELY4cOYJKlSrZBMuievXq2vNCgl4Jii3rE24rFy2Smm8P0mdfjqN1Gr5lP5L7/ZbnLdtKtwBJ8X/Ue7KQixu5EfDrr7/aZf+JiIisyY3y/PnzqwBabi7PnTtXnXfr1q2L48ePu/11TGRkJFq1aoUrV66o975u3TrMmjULJUuWVF3kiNwZ+9gTOYjc5ZY+60JOXoULF1bF3OTkJfOkyB1y6dMufc+7du2KTz/9FIMGDVJ3yfv06aO2efDggerjLXerrVuXn332WdWfXPp9W1rHU2vHjh1qLj+fkrVr16o+57IPcnFhITcGpG97QnKitjxvPbesT7itBMa3bt1C0aJFkRGSGSDHTLII3nnnHZv9fNTvtzxv2Tap7XLkyAGdTmezrZDWA2k92bZtGwYPHpyh/SciIkromWeeUZOFFLKV1mq5gT1//nzMmDHDra9jpEaPnHsly/C5557T1svvInJ3bLEncpCEJz056chd9Y0bNz7y5yRYXrlypaoaL+nhchdaTowW27dvx82bN9XdcAleLZMUqGvbtq1KM5c72mltAZCTsVSHfxTZJ3kf9erVw+TJkxM9L6+RnITPpWXbtJKbA3IBsWXLFnWBUaJEiVT/jozupwT3ku5PRERkb3K+l8KwMoKN3GCW6wqZywgwf//9t9tfx0gGgXQTlMwAec1jx47Z6d0SOT+22BM5SML0bTkZyskuYStvUiSlTu6+//XXX6oCvbQQW0j6mXjxxReT/Xk5YVr/TEokNV760vv4+DyyT5ukv0m/Pkk1t64cL5J7b7IvwtLybWnlT25bCZZl+LuMBPV9+/ZVaX6S3md9Rz81v9+6hV62lSr9CckFh9QlSKo1X4b5k+NJRERkb8OHD1cp6BLYNmnSRAW50qddznv2Pvc443WM9PWXkW0mTpyoWvYtGX79+vVTw/MmVReIyF0wsCdykMuXL6u0bAu5Gy0nQ+v09eSMGzdOFZiRMeNlGDhJV7OkuVta1aWvmrScJ0XS5dJCXlMCVQlYkzqRSlAvY85LgTtJxZcTa0LSb2758uXqfVr3s7cUyqlataqaP/HEE2qc+aQK/cg6uRsvwXFGgnrpzydpepL6l5BlP+R3Sdpfwt9ved7ynlasWKE+S+sLnITvKeHFiIx5T0REZG9y01qGoJVWe2vXr1/P0E1xV7qOsZyb5Zx/6NAhNUxvWFiYurYYNWpUOt4pkWtgKj6RgyxdutRmWfqkyUnRunpsUqQQjKS5y51neSxBdLdu3VTgLaRqrJy8Jf2sVq1aSU6SlpcWMv66SKqau7RYS1Av1eVlf6R1ILkiO/fv38d3331ns15azaVirdy9FxL0SwXf77//3qbQzblz51R6n/TpSw85wcsdewnqpZ9hwkr2FnKRIlX95eJI+iZayFj3UnjI+vdLa79kEMh7sCYXEXIBISmD1uTzlQr8kiJJRERkb3JOSpgx98svv2RKFzBnv46RY1GjRg3MnDlTvZ6k/hO5M7bYEzmIBK4SxEr6+tGjRxEcHKxOQI8q8BIREaFamSW9Tu52S3qd9FOTQjkjR45UlV+lGJzc5Za+adI6LKls0q9bhqCRlDeZf/LJJ2naV8tJWoJb62rxEuhKUC8k7U368MlkIa3vUjVeSBVcea+Scnf37l3V8i4t+FK1V4Jo6zT/8ePHo3bt2uoOvtxdlyr4ckdf7uIHBQUl2jdJu5PA/VGGDBmiWunfeOMNdTdf3ouFXAQ99dRT2vKUKVPUvsrQfAMGDMDVq1fVfkgLvPUNAUkjlL768lnI/ss+S8bCggULMGHChESp+NJyIEWBUju6ABERUVrIeVNuLssNeTlf79u3D9OmTVM33z3hOkaGy503bx46d+6sMgDk2kD28/bt22o/idxaWge+J6KMGTdunESgpn379pk6duxoypkzpylXrlym7t27m65cuWKzbZMmTdQkYmJi1OPChQubIiIibLabNm2aes1Vq1Zp6zZt2mRq3769KV++fCY/Pz/TY489ppb/97//adssWrRI/dzp06dT3O/GjRubnn32WZt1lp9PbpLnrd27d880ZMgQU5EiRUw6nc5UvXp10/Lly5P8fXv37jW1aNHClD17dlPu3LlNnTt3Np08eTLRdjVr1lSvl5JSpUolu5/yXEJr16411atXzxQQEKCO4WuvvZbo8xF6vV59piVLllTvqXz58qaPP/44yX0IDg42FShQwBQVFZXi/hIREaXVrVu3TH369DEVKlRInT8bNWpk2rJli831hJDzfsLzdGqvCZz5Ouaff/5R+/HEE0+YsmXLZgoMDDTVqVPHtHjxYv5nIrfnJf84+uYCkSeRsV6lRVruOKdUZd6ZSAq9pMrJuLPWfeocSVL1pVVc7vAPHDgQzkzS+iVL4ZVXXlHZDURERK7IVa9jiNwd+9gTUapI33JJNU9qGDtH2bx5s7rJIH3nnZ10N5AaA++++66jd4WIiIiI3AwDeyJKFSlCs3DhQlXoTsaSdQbt27fHmTNn0lwM0BHkmEmhIXtXJSYiIiIiYio+ERERERERkQtjiz0RERERERGRC2NgT0REREREROTCGNgTERERERERuTBfR++Aq5DCV5cuXUKuXLlUETEiIiJHkxFrZdhHKWrp7c179RnFcz0REbnquZ6BfSpJUF+iRAl7fT5ERER2c/78eRQvXpxHNIN4riciIlc91zOwTyVpqbcc0Ny5c9vn0yEiIsqAu3fvqpvOlnMUZQzP9URE5Krnegb2qWRJv5egnoE9ERE5E3YRs+9x5LmeiIhc7VzPDnlERERERERELoyBPREREXm0uXPnonLlyqhdu7ajd4WIiChdGNgTERGRRxs4cCCOHTuGPXv2OHpXiIiI0oV97ImIKMuGa4mJiYHRaOQRTyUfHx/4+vqyDz0RuS05JxgMBkfvBpHLn+sZ2BMRUabT6/WIiIjAgwcPeLTTKHv27ChatCh0Oh2PHRG5lfv37+PChQvqxi+RJ8tuh3M9A3siIspUsbGxOH36tLojXaxYMXXSYhX3lMmFrtwQuXbtmjp+5cqVg7c3e9ARkfu01EtQLwFNwYIFeV4gj2Sy47megT0REWUqOWFJcC9jsMoFHKVetmzZ4Ofnh7Nnz6rjGBAQwMNHRG5B0Line truncated
+                            "text/plain": [
+                                "<Figure size 1200x400 with 2 Axes>"
+                            ]
+                        },
+                        "metadata": {},
+                        "output_type": "display_data"
+                    }
+                ],
+            "source": [
+                "x = gpu.bin_centers()\n",
+                "row, col = rows // 2, cols // 2\n",
+                "\n",
+                "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))\n",
+                "ax1.step(x, values_cpu[row, col], where='mid', label='CPU')\n",
+                "ax1.plot(x, values_gpu[row, col], 'rx', ms=5, label='GPU')\n",
+                "ax1.set_xlim(x[values_cpu[row, col].argmax()] + np.array([-100, 100]))\n",
+                "ax1.set(title=f'pixel ({row}, {col})', xlabel='ADU', ylabel='counts')\n",
+                "ax1.legend()\n",
+                "\n",
+                "ax2.step(x, values_cpu.sum(axis=(0, 1)), where='mid', label='CPU')\n",
+                "ax2.plot(x, values_gpu.sum(axis=(0, 1)), 'rx', ms=5, markevery=10, label='GPU (every 10th bin)')\n",
+                "ax2.set_yscale('log')\n",
+                "ax2.set(title='all pixels', xlabel='ADU', ylabel='counts')\n",
+                "ax2.legend()\n",
+                "plt.show()"
+            ]
+        }
+    ],
+    "metadata": {
+        "kernelspec": {
+            "display_name": "Python 3 (ipykernel)",
+            "language": "python",
+            "name": "python3"
+        },
+        "language_info": {
+            "codemirror_mode": {
+                "name": "ipython",
+                "version": 3
+            },
+            "file_extension": ".py",
+            "mimetype": "text/x-python",
+            "name": "python",
+            "nbconvert_exporter": "python",
+            "pygments_lexer": "ipython3",
+            "version": "3.11.15"
+        }
+    },
+    "nbformat": 4,
+    "nbformat_minor": 5
+}

python/tests/test_PixelHistogramCUDA.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: MPL-2.0
+import gc
+import weakref
+
+import numpy as np
+import pytest
+
+from aare import PixelHistogram_u32, PixelHistogramCUDA, _cuda_available
+
+pytestmark = pytest.mark.skipif(
+    not _cuda_available(), reason="aare built without AARE_CUDA")
+
+# Non-square, and integer-valued frames on 10-wide bins: every value is
+# exact in float32 and float64, so the GPU's float axis and the CPU's double
+# axis put each value in the same bin.
+ROWS, COLS, N_FRAMES = 7, 13, 50
+N_BINS, XMIN, XMAX = 120, 3900.0, 5100.0
+
+
+@pytest.fixture
+def frames():
+    rng = np.random.default_rng(12345)
+    base = rng.uniform(4000, 5000, (ROWS, COLS))
+    noise = rng.normal(0, 10, (N_FRAMES, ROWS, COLS))
+    return np.round(base + noise).astype(np.uint16)
+
+
+def make_gpu(dtype, storage_dtype=np.uint32, batch_size=8):
+    return PixelHistogramCUDA(rows=ROWS, cols=COLS, n_bins=N_BINS, xmin=XMIN,
+                              xmax=XMAX, dtype=dtype,
+                              storage_dtype=storage_dtype,
+                              batch_size=batch_size)
+
+
+def cpu_values(frames):
+    cpu = PixelHistogram_u32(rows=ROWS, cols=COLS, n_bins=N_BINS, xmin=XMIN,
+                             xmax=XMAX)
+    for frame in frames:
+        cpu.fill_async(frame.astype(np.float64))
+    return cpu.values()
+
+
+@pytest.mark.parametrize("dtype", [np.uint16, np.float32, np.float64])
+def test_matches_cpu(frames, dtype):
+    data = frames.astype(dtype)
+    gpu = make_gpu(dtype)
+    gpu.fill_histogram(data)
+    np.testing.assert_array_equal(gpu.values(), cpu_values(frames))
+    assert gpu.dtype == np.dtype(dtype)
+
+
+@pytest.mark.parametrize("dtype, axis_dtype", [(np.uint16, np.float32),
+                                               (np.float32, np.float32),
+                                               (np.float64, np.float64)])
+def test_axis_is_floating_point(dtype, axis_dtype):
+    gpu = make_gpu(dtype)
+    edges = gpu.bin_edges()
+    assert edges.dtype == axis_dtype
+    assert edges.shape == (N_BINS + 1,)
+    assert edges[0] == XMIN and edges[-1] == XMAX
+    assert gpu.bin_centers().shape == (N_BINS,)
+
+
+@pytest.mark.parametrize("storage_dtype", [np.uint8, np.uint16, np.uint64,
+                                           np.float32, np.float64])
+def test_storage_dtype(frames, storage_dtype):
+    gpu = make_gpu(np.uint16, storage_dtype=storage_dtype)
+    gpu.fill_histogram(frames)
+    values = gpu.values()
+    assert values.dtype == storage_dtype
+    np.testing.assert_array_equal(values, cpu_values(frames))
+
+
+def test_registered_buffer_outlives_the_caller(frames):
+    gpu = make_gpu(np.uint16)
+    data = frames.copy()
+    gpu.register_input_buffer(data)
+    gpu.fill_histogram(data)
+    del data  # kept alive by the registration
+    np.testing.assert_array_equal(gpu.values(), cpu_values(frames))
+    gpu.unregister_input_buffer()
+
+
+def test_several_fills_accumulate(frames):
+    gpu = make_gpu(np.uint16)
+    gpu.fill_histogram(frames[:20])
+    np.testing.assert_array_equal(gpu.values(), cpu_values(frames[:20]))
+    gpu.fill_histogram(frames[20:])
+    np.testing.assert_array_equal(gpu.values(), cpu_values(frames))
+
+
+def test_read_only_frames_are_accepted(frames):
+    frames.setflags(write=False)
+    gpu = make_gpu(np.uint16)
+    gpu.fill_histogram(frames)
+    np.testing.assert_array_equal(gpu.values(), cpu_values(frames))
+
+
+def test_values_pixel(frames):
+    gpu = make_gpu(np.uint16)
+    gpu.fill_histogram(frames)
+    values = gpu.values()
+    for row, col in [(0, 0), (ROWS - 1, COLS - 1), (3, 6)]:
+        np.testing.assert_array_equal(gpu.values_pixel(row, col),
+                                      values[row, col])
+    with pytest.raises(IndexError):
+        gpu.values_pixel(ROWS, 0)
+
+
+def test_values_into_out(frames):
+    gpu = make_gpu(np.uint16)
+    out = np.empty((ROWS, COLS, N_BINS), dtype=np.uint32)
+    gpu.register_output_buffer(out)
+    gpu.fill_histogram(frames[:20])
+    assert gpu.values(out=out) is out
+    np.testing.assert_array_equal(out, cpu_values(frames[:20]))
+    gpu.fill_histogram(frames[20:])  # the same array again, after more frames
+    gpu.values(out=out)
+    np.testing.assert_array_equal(out, cpu_values(frames))
+    gpu.unregister_output_buffer()
+
+
+def test_registered_out_is_kept_alive():
+    gpu = make_gpu(np.uint16)
+    out = np.empty((ROWS, COLS, N_BINS), dtype=np.uint32)
+    gpu.register_output_buffer(out)
+    ref = weakref.ref(out)
+    del out
+    gc.collect()
+    assert ref() is not None  # still pinned, so it must not be freed
+
+
+def test_bad_out_is_rejected():
+    gpu = make_gpu(np.uint16)
+    good = np.empty((ROWS, COLS, N_BINS), dtype=np.uint32)
+    with pytest.raises(TypeError, match="dtype uint32"):
+        gpu.values(out=good.astype(np.uint16))
+    with pytest.raises(TypeError, match="numpy array"):
+        gpu.values(out=good.tolist())
+    with pytest.raises(ValueError, match="shape"):
+        gpu.values(out=good[:, :, :-1])
+    with pytest.raises(ValueError, match="C-contiguous"):
+        gpu.values(out=np.empty((N_BINS, COLS, ROWS), np.uint32).transpose())
+    good.setflags(write=False)
+    with pytest.raises(ValueError, match="writeable"):
+        gpu.values(out=good)
+    with pytest.raises(ValueError, match="shape"):
+        gpu.register_output_buffer(np.empty((ROWS, COLS), np.uint32))
+
+
+def test_wrong_dtype_is_not_converted(frames):
+    gpu = make_gpu(np.uint16)
+    with pytest.raises(TypeError, match="dtype=float64"):
+        gpu.fill_histogram(frames.astype(np.float64))
+    with pytest.raises(TypeError):
+        gpu.register_input_buffer(frames.astype(np.float32))
+
+
+def test_bad_layout_is_rejected(frames):
+    gpu = make_gpu(np.uint16)
+    with pytest.raises(ValueError, match="C-contiguous"):
+        gpu.fill_histogram(frames[:, :, ::-1])
+    with pytest.raises(ValueError, match="3D"):
+        gpu.fill_histogram(frames[0])
+
+
+def test_factory_rejects_unsupported_dtypes():
+    with pytest.raises(ValueError, match="Unsupported dtype"):
+        make_gpu(np.int32)
+    with pytest.raises(ValueError, match="Unsupported storage_dtype"):
+        make_gpu(np.uint16, storage_dtype=np.int8)
+
+
+def test_constructor_validates_arguments():
+    with pytest.raises(ValueError):
+        PixelHistogramCUDA(rows=ROWS, cols=COLS, n_bins=N_BINS, xmin=XMAX,
+                           xmax=XMIN, dtype=np.uint16)
+    with pytest.raises(ValueError):
+        PixelHistogramCUDA(rows=ROWS, cols=COLS, n_bins=N_BINS, xmin=XMIN,
+                           xmax=XMAX, dtype=np.uint16, batch_size=0)