aare/python/tests/ClusterFinderCUDA.ipynb

In [1]:

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

from pathlib import Path
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np
import boost_histogram as bh
import time

from aare import File, ClusterFinder, ClusterFinderMT, ClusterCollector, ClusterFinderCUDA

In [2]:

# Helpers
N_BINS = 200
def make_hist(clusters):
    h = bh.Histogram(bh.axis.Regular(N_BINS, -2, 4000))
    h.fill(clusters.sum())
    return h

def make_hist_from_batch(result_list):
    h = bh.Histogram(bh.axis.Regular(N_BINS, -2, 4000))
    energies = [np.asarray(cv.sum()).ravel() for cv in result_list if cv.size > 0]
    if energies:
        h.fill(np.concatenate(energies))
    return h

In [3]:

base = Path('/mnt/sls_det_storage/matterhorn_data/aare_test_data/')
f = File(base / 'Moench03new/cu_half_speed_master_4.json')

n_frames_pd = 1000
N           = 88999 #88999
cluster_size = (9, 9)
rows = f.rows
cols = f.cols
image_size   = (rows, cols)
capacity     = 50_000 #3_000_000

print(f'Image size:       {image_size}')
print(f'Pedestal frames:  {n_frames_pd}')
print(f'Data frames:      {N}')

Image size:       (400, 400)
Pedestal frames:  1000
Data frames:      88999

In [4]:

f.total_frames

Out [4]:

500000

Pedestal (both finders trained on identical frames)

• Modify the boolean SERIAL to choose between the sequential CPU version (ClusterFinder) and its multi-threaded homologue (ClusterFinderMT)

In [5]:

SERIAL = True

In [6]:

if(SERIAL):
    cf_cpu  = ClusterFinder(image_size, cluster_size, capacity=capacity)
else:
    cf_cpu  = ClusterFinderMT(image_size, cluster_size, capacity=capacity, n_threads=48)
    sink = ClusterCollector(cf_cpu)

In [7]:

# Runs the destructor under the hood in case cf_cuda has already been constructed
# del cf_cuda
cf_cuda = None

In [8]:

N_STREAMS  = 5
cf_cuda = ClusterFinderCUDA(image_size, 
                            cluster_size, 
                            n_sigma=7, 
                            max_clusters_per_frame=1500,
                            n_streams=N_STREAMS) 

In [9]:

t0 = time.perf_counter()
for _ in range(n_frames_pd):
    img = f.read_frame()
    cf_cpu.push_pedestal_frame(img.copy())
    cf_cuda.push_pedestal_frame(img.copy())
print(f'Pedestal ({n_frames_pd} frames): {time.perf_counter() - t0:.3f}s')

Pedestal (1000 frames): 0.499s

Read all data frames into memory (I/O out of the timing loop)

In [10]:

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

Reading 88999 frames:        51.133s  (1741 FPS, 531.166 GB/s)

CPU clustering

In [11]:

from tqdm import tqdm

In [12]:

t0 = time.perf_counter()
for frame in tqdm(data):
    cf_cpu.find_clusters(frame)
t_cpu = time.perf_counter() - t0

if(SERIAL):
    clusters_cpu = cf_cpu.steal_clusters(realloc_same_capacity=False)
    n_clusters_cpu = clusters_cpu.size
    
    hist_cpu  = make_hist(clusters_cpu)
else:
    cf_cpu.stop()
    sink.stop()
    
    clusters_cpu = sink.steal_clusters() #cf_cpu.steal_clusters(realloc_same_capacity=False)
    
    hist_cpu = bh.Histogram(bh.axis.Regular(N_BINS, -2, 4000))
    n_clusters_cpu = 0
    for cv in clusters_cpu:
        hist_cpu.fill(cv.sum())
        n_clusters_cpu += cv.size
        
print(f'CPU clustering:          {t_cpu:.3f}s ({N/t_cpu:.0f} FPS, '
      f'{n_clusters_cpu} clusters, {n_clusters_cpu/N:.2f}/frame)')

100%|████████████████████████████████████| 88999/88999 [13:59<00:00, 106.07it/s]

CPU clustering:          839.049s (106 FPS, 90799856 clusters, 1020.23/frame)

CUDA clustering

In [13]:

BATCHED = True

In [29]:

if(BATCHED):
    BATCH_SIZE = 2000

    # Before warmup, pin the fixed size buffer
    batch_buffer = np.empty((BATCH_SIZE, rows, cols), dtype=np.uint16)
    cf_cuda.register_input_buffer(batch_buffer)       # fixed ~640 MB

    # Warmup: first kernel launch pays CUDA context + pedestal H2D upload cost
    _ = cf_cuda.find_clusters_batched(data[0:BATCH_SIZE], first_frame=0)
    
    clusters_cuda_per_frame = []

    cf_cuda.reset_timers()
    t0 = time.perf_counter()
    for start in range(0, N, BATCH_SIZE):
        stop = min(start + BATCH_SIZE, N)
        batch_buffer[:stop-start] = data[start:stop]  # CPU memcpy into pinned buf
        clusters_cuda_per_frame.extend(
            cf_cuda.find_clusters_batched(batch_buffer[:stop-start], first_frame=start)
        )
    t_cuda = time.perf_counter() - t0

    cf_cuda.unregister_input_buffer()  # release when done with this dataset

    kernel_ms = cf_cuda.avg_kernel_time_ms()
    
    n_clusters_cuda = sum(cv.size for cv in clusters_cuda_per_frame)

    hist_cuda = make_hist_from_batch(clusters_cuda_per_frame)
    
else:   
    # Simpler: (non-batched) per-frame run on non-pinned data
    cf_cuda.find_clusters(data[0])
    _ = cf_cuda.steal_clusters(realloc_same_capacity=False)

    cf_cuda.reset_timers()
    t0 = time.perf_counter()

    n_clusters_cuda = 0
    hist_cuda = None

    # steal the clusters as we go rather than at the end of the  dataset 
    # which might trigger an std::bad_alloc...
    for idx, frame in enumerate(data):
        cf_cuda.find_clusters(frame)
        clusters_frame = cf_cuda.steal_clusters(realloc_same_capacity=True)

        n_clusters_cuda += clusters_frame.size

        h = make_hist(clusters_frame)
        hist_cuda = h if hist_cuda is None else hist_cuda + h
        
    t_cuda = time.perf_counter() - t0
    kernel_ms = cf_cuda.avg_kernel_time_ms()

In [30]:

cluster_size

Out [30]:

(9, 9)

In [31]:

print(f'CPU clustering:          {t_cpu:.3f}s ({N/t_cpu:.0f} FPS, '
      f'{n_clusters_cpu} clusters, {n_clusters_cpu/N:.2f}/frame)')

CPU clustering:          839.049s (106 FPS, 90799856 clusters, 1020.23/frame)

In [32]:

print(f'CUDA clustering:          {t_cuda:.3f}s  ({N/t_cuda:.0f} FPS, '
      f'{n_clusters_cuda} clusters, {n_clusters_cuda/N:.2f}/frame)')
print(f'  Kernel only:            {kernel_ms:.3f} ms/frame')
print(f'  PCIe + overhead:        {t_cuda*1000/N - kernel_ms:.3f} ms/frame')
print(f'Speedup (CPU / CUDA):     {t_cpu / t_cuda:.2f}×')

CUDA clustering:          7.039s  (12644 FPS, 89991186 clusters, 1011.15/frame)
  Kernel only:            0.063 ms/frame
  PCIe + overhead:        0.016 ms/frame
Speedup (CPU / CUDA):     119.20×

Agreement check:

• Cluster counts should match closely.
• However, as the CUDA CF updates the pedestal once per frame rather than per-pixel, a small divergence after the first few frames is expected.

In [18]:

diff = abs(n_clusters_cpu - n_clusters_cuda)
rel  = diff / max(n_clusters_cpu, 1)
print(f'Cluster count diff:       {diff} ({rel:.2%})')

Cluster count diff:       2716962 (2.99%)

Plots

In [19]:

print(len(hist_cpu.values()), len(hist_cpu.axes[0].edges))
print(len(hist_cuda.values()), len(hist_cuda.axes[0].edges))

200 201
200 201

In [20]:

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

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

ax_spec.stairs(cpu_vals, edges, label=f'CPU  ({n_clusters_cpu} clusters)')
ax_spec.stairs(cuda_vals, edges, label=f'CUDA ({n_clusters_cuda} clusters)', linestyle='--')
ax_spec.set_ylabel('Counts')
ax_spec.set_title('Cluster energy spectrum: CPU vs CUDA')
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='k')
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.5, 2.0)
ax_ratio.grid(alpha=0.3)

plt.tight_layout()
plt.show()

In [ ]: