aare/python/tests/ClusterFinderCUDA.ipynb

ClusterFinder — CPU vs CUDA¶

Two finders over the same frames and the same pedestal:

ClusterFinderMT + ClusterCollector	multithreaded CPU baseline
ClusterFinderCUDA.find_clusters_batched	batched across CUDA streams, over a pinned input buffer, so H2D / kernel / D2H overlap

Part 1 times them. Part 2 compares what they found.

In [1]:

import sys; sys.path.append('/home/ferjao_k/aare/build')

import time
from pathlib import Path

import boost_histogram as bh
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm

from aare import (File, ClusterFinderMT, ClusterCollector, ClusterFinderCUDA)

Configuration¶

In [2]:

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

n_frames_pd  = 1000                  # pedestal frames, taken from the head of the file
N            = 50000                 # data frames, read after those
cluster_size = (3, 3)
rows, cols   = f.rows, f.cols
image_size   = (rows, cols)

N_SIGMA    = 5
N_STREAMS  = 4                       # H2D / kernel / D2H overlap
BATCH_SIZE = 2000                    # frames per find_clusters_batched call

# CPU workers. ClusterFinderMT gives each thread its own ClusterFinder with its
# own pedestal, and data frames are handed out round-robin -- so at high thread
# counts each pedestal tracks on only 1/n of the frames while the CUDA finder's
# single pedestal sees all of them, and the cluster counts drift apart. Keep it
# low when comparing counts; raise it to measure CPU throughput.
N_THREADS = 24

# Fixed per-frame output slot the kernel writes into, copied WHOLE on every D2H.
# Too low does not error -- it TRUNCATES. The CUDA cell prints the peak actually
# reached, which is how you size this for a new dataset.
CAP = 3000

capacity = 10_000                    # initial ClusterVector capacity, CPU side
N_BINS   = 200
E_MIN, E_MAX = -2, 4000              # spectrum range, ADU

print(f'Image size:      {image_size}   cluster {cluster_size}')
print(f'Pedestal frames: {n_frames_pd}')
print(f'Data frames:     {N}  of {f.total_frames} in file')

Image size:      (400, 400)   cluster (3, 3)
Pedestal frames: 1000
Data frames:     50000  of 100000 in file

Build the two finders¶

In [3]:

cf_cpu = ClusterFinderMT(image_size, cluster_size, n_sigma=N_SIGMA,
                         capacity=capacity, n_threads=N_THREADS)
sink = ClusterCollector(cf_cpu)

cf_cuda = ClusterFinderCUDA(image_size, cluster_size, n_sigma=N_SIGMA,
                            max_clusters_per_frame=CAP, n_streams=N_STREAMS)

Pedestal¶

Identical frames into both, so the implementation is the only thing that differs.

In [4]:

pd.seek(0)
t0 = time.perf_counter()
for _ in range(n_frames_pd):
    pd_img = pd.read_frame()
    cf_cpu.push_pedestal_frame(pd_img.copy())
    cf_cuda.push_pedestal_frame(pd_img.copy())
print(f'Pedestal ({n_frames_pd} frames): {time.perf_counter() - t0:.2f}s')

Pedestal (1000 frames): 2.01s

Read the data frames¶

Kept out of both timing loops: the two finders run over the same in-memory array.

In [5]:

f.seek(0)
t0 = time.perf_counter()
data = f.read_n(N)
t_io = time.perf_counter() - t0
gb = f.bytes_per_frame * N / 1e9
print(f'Reading {N} frames: {t_io:.2f}s  ({N / t_io:,.0f} FPS, '
      f'{gb:.2f} GB at {gb / t_io:.2f} GB/s)')

Reading 50000 frames: 31.25s  (1,600 FPS, 16.00 GB at 0.51 GB/s)

1 — CPU, multithreaded¶

In [6]:

t0 = time.perf_counter()
# Pass the frame number: find_clusters(frame) defaults it to 0, and every
# resulting cluster is then stamped frame 0. Part 2 groups by it.
for i, frame in enumerate(tqdm(data)):
    cf_cpu.find_clusters(frame, i)
t_cpu = time.perf_counter() - t0

cf_cpu.stop(); sink.stop()
clusters_cpu = sink.steal_clusters()

hist_cpu = bh.Histogram(bh.axis.Regular(N_BINS, E_MIN, E_MAX))
n_clusters_cpu = 0
for cv in clusters_cpu:
    hist_cpu.fill(cv.sum())
    n_clusters_cpu += cv.size

print(f'CPU MT ({N_THREADS} threads): {t_cpu:.2f}s   {N / t_cpu:,.0f} FPS   '
      f'{n_clusters_cpu:,} clusters ({n_clusters_cpu / N:.1f}/frame)')

100%|███████████████████████████████████████| 50000/50000 [00:12<00:00, 3977.38it/s]

CPU MT (24 threads): 12.58s   3,976 FPS   116,015,414 clusters (2320.3/frame)

2 — CUDA, batched + multi-streamed¶

register_input_buffer page-locks the input once so H2D runs at full DMA bandwidth instead of through the driver's staging buffer. It is paid once, outside the loop, and released when the dataset is done.

In [7]:

t0 = time.perf_counter()
cf_cuda.register_input_buffer(data)
t_pin = time.perf_counter() - t0

clusters_cuda = []
t0 = time.perf_counter()
for start in range(0, N, BATCH_SIZE):
    stop = min(start + BATCH_SIZE, N)
    clusters_cuda.extend(
        cf_cuda.find_clusters_batched(data[start:stop], first_frame=start))
t_cuda = time.perf_counter() - t0

cf_cuda.unregister_input_buffer()

hist_cuda = bh.Histogram(bh.axis.Regular(N_BINS, E_MIN, E_MAX))
energies = [np.asarray(cv.sum()).ravel() for cv in clusters_cuda if cv.size > 0]
if energies:
    hist_cuda.fill(np.concatenate(energies))
n_clusters_cuda = sum(cv.size for cv in clusters_cuda)

print(f'CUDA batched: {t_cuda:.2f}s   {N / t_cuda:,.0f} FPS   '
      f'{n_clusters_cuda:,} clusters ({n_clusters_cuda / N:.1f}/frame)')
print(f'  input pinned in {t_pin:.2f}s, once, outside the loop')
print(f'  vs CPU: {t_cpu / t_cuda:.2f}x')

# One ClusterVector per frame, so this IS the per-frame maximum.
peak = max((cv.size for cv in clusters_cuda), default=0)
print(f'  peak clusters in one frame: {peak:,} of CAP {CAP:,}')
if peak >= CAP:
    print('  !! CAP REACHED -- clusters were dropped; raise CAP and re-run.')

CUDA batched: 1.87s   26,720 FPS   116,013,348 clusters (2320.3/frame)
  input pinned in 0.53s, once, outside the loop
  vs CPU: 6.72x
  peak clusters in one frame: 2,524 of CAP 3,000

Summary¶

In [8]:

print(f'{"":<22} {"s":>7} {"FPS":>10} {"us/frame":>10} {"clusters":>13}')
for name, t, n in (('CPU MT (%d thr)' % N_THREADS, t_cpu, n_clusters_cpu),
                   ('CUDA batched', t_cuda, n_clusters_cuda)):
    print(f'{name:<22} {t:7.2f} {N / t:10,.0f} {t * 1e6 / N:10.1f} {n:>13,}')
print(f'\nspeedup: {t_cpu / t_cuda:.2f}x')

                             s        FPS   us/frame      clusters
CPU MT (24 thr)          12.58      3,976      251.5   116,015,414
CUDA batched              1.87     26,720       37.4   116,013,348

speedup: 6.72x

Part 2 — Where the two disagree¶

The two implementations are not bit-identical by construction: the CUDA finder updates its pedestal once per frame, the CPU finder per pixel, so once the pedestals drift apart a few pixels land on opposite sides of the threshold.

Three views, cheapest first: the spectra, the cluster sets, and where in the run the difference accumulates.

Spectra¶

Both curves should lie on top of each other; the ratio panel is the sensitive view.

In [9]:

fig, (ax_spec, ax_ratio) = plt.subplots(
    2, 1, figsize=(9, 6.5), sharex=True, gridspec_kw={'height_ratios': [3, 1]})

edges = hist_cpu.axes[0].edges
cpu_vals, cuda_vals = hist_cpu.values(), hist_cuda.values()

ax_spec.stairs(cpu_vals, edges, label=f'CPU MT  ({n_clusters_cpu:,})',
               color='C0', linewidth=1.4)
ax_spec.stairs(cuda_vals, edges, label=f'CUDA  ({n_clusters_cuda:,})',
               color='C1', linestyle='--', linewidth=1.4)
ax_spec.set_ylabel('Counts')
ax_spec.set_title(f'Cluster energy spectrum — {N:,} frames')
ax_spec.legend()
ax_spec.grid(alpha=0.2)

with np.errstate(divide='ignore', invalid='ignore'):
    ratio = np.where(cpu_vals > 0, cuda_vals / cpu_vals, np.nan)
ax_ratio.stairs(ratio, edges, color='C1', linewidth=1.2)
ax_ratio.axhline(1.0, color='gray', linewidth=0.5)
ax_ratio.set_ylabel('CUDA / CPU')
ax_ratio.set_xlabel('Energy [ADU]')
ax_ratio.set_ylim(0.9, 1.1)
ax_ratio.grid(alpha=0.3)

plt.tight_layout()
plt.show()

Cluster sets¶

Every cluster is keyed by (frame, x, y), so the two results can be differenced directly — no per-frame Python loop.

A cluster whose centre moved by one pixel was relabelled, not gained or lost, so those are counted separately from genuine adds and drops.

In [10]:

def cluster_keys(vectors):
    """One int64 key per cluster: frame * rows*cols + y * cols + x.

    Centres are unique within a frame, so keys are globally unique and the sets
    below need only a sort, never a np.unique.
    """
    n = len(vectors)
    sizes = np.fromiter((cv.size for cv in vectors), np.int64, n)
    frames = np.repeat(
        np.fromiter((cv.frame_number for cv in vectors), np.int64, n), sizes)
    xy = [np.asarray(cv) for cv in vectors if cv.size]
    if not xy:
        return np.empty(0, np.int64), frames
    x = np.concatenate([v['x'] for v in xy]).astype(np.int64)
    y = np.concatenate([v['y'] for v in xy]).astype(np.int64)
    return frames * (rows * cols) + y * cols + x, frames


t0 = time.perf_counter()
k_cpu, fr_cpu = cluster_keys(clusters_cpu)
k_cuda, fr_cuda = cluster_keys(clusters_cuda)

# Sort ONCE. np.setdiff1d and np.isin each re-sort their argument on every call,
# and the shift check below calls one eight times per direction -- about twenty
# sorts of a 27-million-element array. Two sorts plus binary search instead.
s_cpu = np.sort(k_cpu)
s_cuda = np.sort(k_cuda)


def present_in(vals, sorted_keys):
    """Which of `vals` appear in an already-sorted key array."""
    if sorted_keys.size == 0:
        return np.zeros(vals.size, dtype=bool)
    i = np.searchsorted(sorted_keys, vals)
    np.clip(i, 0, sorted_keys.size - 1, out=i)
    return sorted_keys[i] == vals


cpu_only = s_cpu[~present_in(s_cpu, s_cuda)]
cuda_only = s_cuda[~present_in(s_cuda, s_cpu)]


def drop_shifted(only, sorted_other):
    """Drop mismatches that have a counterpart within 1 px in the other set."""
    keep = np.ones(only.size, dtype=bool)
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            if dx or dy:
                keep &= ~present_in(only + dy * cols + dx, sorted_other)
    return only[keep]


cpu_real = drop_shifted(cpu_only, s_cuda)
cuda_real = drop_shifted(cuda_only, s_cpu)
touched = np.unique(np.concatenate([cpu_real, cuda_real]) // (rows * cols)).size

print(f'CPU  {k_cpu.size:>12,} clusters')
print(f'CUDA {k_cuda.size:>12,} clusters      net {k_cuda.size - k_cpu.size:+,}'
      f'  ({abs(k_cuda.size - k_cpu.size) / k_cpu.size:.4%})')
print()
print(f'exact-position mismatches   CPU-only {cpu_only.size:>8,}   '
      f'CUDA-only {cuda_only.size:>8,}')
print(f'  of which shifted by 1 px  {cpu_only.size - cpu_real.size:>17,}   '
      f'{cuda_only.size - cuda_real.size:>17,}')
print(f'  genuine add / drop        {cpu_real.size:>17,}   {cuda_real.size:>17,}')
print(f'\nframes touched by a genuine difference: {touched:,} of {N:,}')
print(f'[{time.perf_counter() - t0:.1f}s]')

CPU   116,015,414 clusters
CUDA  116,013,348 clusters      net -2,066  (0.0018%)

exact-position mismatches   CPU-only   67,148   CUDA-only   65,082
  of which shifted by 1 px             52,084              51,984
  genuine add / drop                   15,064              13,098

frames touched by a genuine difference: 20,426 of 50,000
[34.1s]

Where in the run it accumulates¶

Cumulative CUDA − CPU against frame index. A curve that rises and then flattens means the two pedestals settled into agreement and the difference is a start-up transient; a steadily sloping line would mean a persistent per-frame bias.

In [11]:

cnt_cpu = np.bincount(fr_cpu, minlength=N)
cnt_cuda = np.bincount(fr_cuda, minlength=N)
cum = np.cumsum(cnt_cuda.astype(np.int64) - cnt_cpu.astype(np.int64))

fig, ax = plt.subplots(figsize=(9, 3.6))
ax.plot(cum, color='C1', linewidth=1.2)
ax.axhline(0, color='gray', linewidth=0.5)
ax.set_xlabel('Frame')
ax.set_ylabel('cumulative CUDA - CPU')
ax.set_title(f'Cluster-count difference accumulating over {N:,} frames '
             f'(final {cum[-1]:+,})')
ax.grid(alpha=0.25)
plt.tight_layout()
plt.show()

half = N // 2
print(f'first half: {cum[half - 1]:+,}      second half: {cum[-1] - cum[half - 1]:+,}')

first half: -468      second half: -1,598