Optimize CUDA cluster finder transfers and kernel hot path

- Use per-stream pinned host staging buffers for truly async CUDA transfers.
- Avoid reserving full device capacity per result frame.
- Reduce kernel work by delaying cluster payload construction.
- Use squared comparisons and removing per-pixel sqrtf() ops.

include/aare/ClusterFinderCUDA.hpp

@@ -3,9 +3,12 @@
 #include "aare/ClusterFinder.hpp"
 #include "aare/clusterfinder_kernel.cuh"
 #include "aare/utils/cuda_check.cuh"
+#include <algorithm>
 #include <cmath>
 #include <cstdint>
 #include <cstdio>
+#include <cstring>
+#include <limits>
 #include <stdexcept>
 
 namespace aare {
@@ -21,6 +24,12 @@ struct StreamContext {
     ClusterType *d_clusters = nullptr;
     uint32_t *d_cluster_count = nullptr;
 
+    // Pinned host staging buffers.  These make cudaMemcpyAsync real async DMA
+    // transfers even when the caller's NDView points to pageable memory.
+    FRAME_TYPE *h_frame = nullptr;
+    uint32_t *h_cluster_count = nullptr;
+    ClusterType *h_clusters = nullptr;
+
     cudaEvent_t kernel_start = nullptr;
     cudaEvent_t kernel_stop = nullptr;
 };
@@ -44,6 +53,9 @@ class ClusterFinderCUDA {
     int n_streams;
     size_t m_capacity;
 
+    size_t m_image_bytes;
+    size_t m_cluster_bytes;
+
     COMPUTE_TYPE m_nSigma;
     Pedestal<PEDESTAL_TYPE> m_pedestal;
     ClusterVector<ClusterType> m_clusters;
@@ -77,6 +89,22 @@ class ClusterFinderCUDA {
           m_image_size(nrows * ncols), n_streams(n_streams_),
           m_capacity(capacity), m_nSigma(nSigma),
           m_pedestal(shape_[0], shape_[1]), m_clusters(capacity) {
+        if (n_streams_ <= 0) {
+            throw std::invalid_argument(
+                "ClusterFinderCUDA: n_streams must be > 0");
+        }
+
+        if (capacity >
+            static_cast<size_t>(std::numeric_limits<uint32_t>::max())) {
+            throw std::invalid_argument(
+                "ClusterFinderCUDA: capacity must fit in uint32_t");
+        }
+
+        if (capacity == 0) {
+            throw std::invalid_argument(
+                "ClusterFinderCUDA: capacity must be > 0");
+        }
+
         // Grid/Block dimensions
         block = dim3(BLOCK_X, BLOCK_Y);
         grid = dim3((static_cast<unsigned int>(ncols) + BLOCK_X - 1) / BLOCK_X,
@@ -89,28 +117,44 @@ class ClusterFinderCUDA {
         shmem_bytes = (BLOCK_X + 2 * col_radius) * (BLOCK_Y + 2 * row_radius) *
                       sizeof(COMPUTE_TYPE);
 
+        m_image_bytes = m_image_size * sizeof(FRAME_TYPE);
+        m_cluster_bytes = m_capacity * sizeof(ClusterType);
+
         v_sc.resize(n_streams);
         for (int k = 0; k < n_streams; ++k) {
             auto &sc = v_sc[k];
-            CUDA_CHECK(cudaStreamCreate(&sc.stream));
+            CUDA_CHECK(
+                cudaStreamCreateWithFlags(&sc.stream, cudaStreamNonBlocking));
             CUDA_CHECK(cudaEventCreate(&sc.kernel_start));
             CUDA_CHECK(cudaEventCreate(&sc.kernel_stop));
-            CUDA_CHECK(
-                cudaMalloc(&sc.d_frame, m_image_size * sizeof(FRAME_TYPE)));
+            CUDA_CHECK(cudaMalloc(&sc.d_frame, m_image_bytes));
             CUDA_CHECK(cudaMalloc(&sc.d_pd_mean,
                                   m_image_size * sizeof(PEDESTAL_TYPE)));
             CUDA_CHECK(
                 cudaMalloc(&sc.d_pd_sum, m_image_size * sizeof(PEDESTAL_TYPE)));
             CUDA_CHECK(cudaMalloc(&sc.d_pd_sum2,
                                   m_image_size * sizeof(PEDESTAL_TYPE)));
-            CUDA_CHECK(
-                cudaMalloc(&sc.d_clusters, capacity * sizeof(ClusterType)));
+            CUDA_CHECK(cudaMalloc(&sc.d_clusters, m_cluster_bytes));
             CUDA_CHECK(cudaMalloc(&sc.d_cluster_count, sizeof(uint32_t)));
+
+            CUDA_CHECK(cudaMallocHost(reinterpret_cast<void **>(&sc.h_frame),
+                                      m_image_bytes));
+            CUDA_CHECK(
+                cudaMallocHost(reinterpret_cast<void **>(&sc.h_cluster_count),
+                               sizeof(uint32_t)));
+            if (m_cluster_bytes > 0) {
+                CUDA_CHECK(
+                    cudaMallocHost(reinterpret_cast<void **>(&sc.h_clusters),
+                                   m_cluster_bytes));
+            }
         }
     }
 
     ~ClusterFinderCUDA() {
         for (auto &sc : v_sc) {
+            if (sc.stream)
+                cudaStreamSynchronize(sc.stream);
+
             if (sc.d_frame)
                 cudaFree(sc.d_frame);
             if (sc.d_pd_mean)
@@ -123,12 +167,20 @@ class ClusterFinderCUDA {
                 cudaFree(sc.d_clusters);
             if (sc.d_cluster_count)
                 cudaFree(sc.d_cluster_count);
-            if (sc.stream)
-                cudaStreamDestroy(sc.stream);
+
+            if (sc.h_frame)
+                cudaFreeHost(sc.h_frame);
+            if (sc.h_clusters)
+                cudaFreeHost(sc.h_clusters);
+            if (sc.h_cluster_count)
+                cudaFreeHost(sc.h_cluster_count);
+
             if (sc.kernel_start)
                 cudaEventDestroy(sc.kernel_start);
             if (sc.kernel_stop)
                 cudaEventDestroy(sc.kernel_stop);
+            if (sc.stream)
+                cudaStreamDestroy(sc.stream);
         }
     }
 
@@ -179,16 +231,18 @@ class ClusterFinderCUDA {
         }
 
         auto &sc = v_sc[0];
-        const size_t image_bytes = m_image_size * sizeof(FRAME_TYPE);
         const uint32_t n_pd_samples =
             static_cast<uint32_t>(m_pedestal.n_samples());
 
+        // First, CPU copies frame into a reusable pinned buffer
+        std::memcpy(sc.h_frame, frame.data(), m_image_bytes);
+
         // Reset cluster counter
         CUDA_CHECK(cudaMemsetAsync(sc.d_cluster_count, 0, sizeof(uint32_t),
                                    sc.stream));
 
         // Upload frame
-        CUDA_CHECK(cudaMemcpyAsync(sc.d_frame, frame.data(), image_bytes,
+        CUDA_CHECK(cudaMemcpyAsync(sc.d_frame, sc.h_frame, m_image_bytes,
                                    cudaMemcpyHostToDevice, sc.stream));
 
         // Timed Kernel launch
@@ -202,9 +256,8 @@ class ClusterFinderCUDA {
         CUDA_CHECK(cudaEventRecord(sc.kernel_stop, sc.stream));
         CUDA_CHECK(cudaGetLastError());
 
-        // Read back cluster count
-        uint32_t n_found = 0;
-        CUDA_CHECK(cudaMemcpyAsync(&n_found, sc.d_cluster_count,
+        // Read back cluster count into pinned buffer
+        CUDA_CHECK(cudaMemcpyAsync(sc.h_cluster_count, sc.d_cluster_count,
                                    sizeof(uint32_t), cudaMemcpyDeviceToHost,
                                    sc.stream));
 
@@ -212,9 +265,11 @@ class ClusterFinderCUDA {
         // clusters
         CUDA_CHECK(cudaStreamSynchronize(sc.stream));
 
+        record_kernel_time(sc);
+
         // Clamp to max in case of overflow
-        if (n_found > m_capacity)
-            n_found = m_capacity;
+        uint32_t n_found = *sc.h_cluster_count;
+        n_found = std::min(n_found, static_cast<uint32_t>(m_capacity));
 
         // Read back clusters
         m_clusters.set_frame_number(frame_number);
@@ -239,19 +294,16 @@ class ClusterFinderCUDA {
         }
 
         const size_t n_frames = frames.shape(0);
-        const size_t image_bytes = m_image_size * sizeof(FRAME_TYPE);
         const uint32_t n_pd_samples =
             static_cast<uint32_t>(m_pedestal.n_samples());
 
         std::vector<ClusterVector<ClusterType>> results;
         results.reserve(n_frames);
         for (size_t i = 0; i < n_frames; ++i) {
-            results.emplace_back(m_capacity);
+            results.emplace_back();
             results.back().set_frame_number(first_frame + i);
         }
 
-        std::vector<uint32_t> host_counts(n_streams, 0);
-
         const size_t n_rounds = (n_frames + n_streams - 1) / n_streams;
         for (size_t round = 0; round < n_rounds; ++round) {
 
@@ -266,11 +318,13 @@ class ClusterFinderCUDA {
                 const FRAME_TYPE *h_src =
                     frames.data() + frame_idx * m_image_size;
 
+                std::memcpy(sc_k.h_frame, h_src, m_image_bytes);
+
                 CUDA_CHECK(cudaMemsetAsync(sc_k.d_cluster_count, 0,
                                            sizeof(uint32_t), sc_k.stream));
-                CUDA_CHECK(cudaMemcpyAsync(sc_k.d_frame, h_src, image_bytes,
-                                           cudaMemcpyHostToDevice,
-                                           sc_k.stream));
+                CUDA_CHECK(
+                    cudaMemcpyAsync(sc_k.d_frame, sc_k.h_frame, m_image_bytes,
+                                    cudaMemcpyHostToDevice, sc_k.stream));
 
                 CUDA_CHECK(cudaEventRecord(sc_k.kernel_start, sc_k.stream));
                 device::find_clusters_in_single_frame<ClusterType, FRAME_TYPE,
@@ -278,10 +332,15 @@ class ClusterFinderCUDA {
                     <<<grid, block, shmem_bytes, sc_k.stream>>>(
                         sc_k.d_frame, sc_k.d_pd_mean, sc_k.d_pd_sum,
                         sc_k.d_pd_sum2, n_pd_samples, m_nSigma, nrows, ncols,
-                        sc_k.d_clusters, sc_k.d_cluster_count, m_capacity);
+                        sc_k.d_clusters, sc_k.d_cluster_count,
+                        static_cast<uint32_t>(m_capacity));
                 CUDA_CHECK(cudaEventRecord(sc_k.kernel_stop, sc_k.stream));
-
                 CUDA_CHECK(cudaGetLastError());
+
+                // Queue count D2H immediately after the kernel
+                CUDA_CHECK(cudaMemcpyAsync(
+                    sc_k.h_cluster_count, sc_k.d_cluster_count,
+                    sizeof(uint32_t), cudaMemcpyDeviceToHost, sc_k.stream));
             }
 
             // Drain phase: fan in results from all streams
@@ -292,14 +351,14 @@ class ClusterFinderCUDA {
 
                 auto &sc_k = v_sc[k];
 
-                CUDA_CHECK(cudaMemcpyAsync(
-                    &host_counts[k], sc_k.d_cluster_count, sizeof(uint32_t),
-                    cudaMemcpyDeviceToHost, sc_k.stream));
+                // Wait for memset -> H2D -> kernel -> count D2H
                 CUDA_CHECK(cudaStreamSynchronize(sc_k.stream));
 
-                uint32_t n_found = host_counts[k];
-                if (n_found > m_capacity)
-                    n_found = m_capacity;
+                record_kernel_time(sc_k);
+
+                uint32_t n_found = *sc_k.h_cluster_count;
+                n_found = std::min<uint32_t>(n_found,
+                                             static_cast<uint32_t>(m_capacity));
 
                 if (n_found > 0) {
                     append_device_clusters_to(results[frame_idx], sc_k,
@@ -354,20 +413,22 @@ class ClusterFinderCUDA {
      */
     void append_device_clusters_to(ClusterVector<ClusterType> &cv, SC &sc,
                                    uint32_t n_found) {
-        std::vector<ClusterType> staging(n_found);
-        CUDA_CHECK(cudaMemcpyAsync(staging.data(), sc.d_clusters,
+
+        CUDA_CHECK(cudaMemcpyAsync(sc.h_clusters, sc.d_clusters,
                                    n_found * sizeof(ClusterType),
                                    cudaMemcpyDeviceToHost, sc.stream));
+
         CUDA_CHECK(cudaStreamSynchronize(sc.stream));
 
-        // Record the total time of all the kernel launches compute time
+        for (uint32_t i = 0; i < n_found; ++i)
+            cv.push_back(sc.h_clusters[i]);
+    }
+
+    void record_kernel_time(SC &sc) {
         float ms = 0.0f;
         CUDA_CHECK(cudaEventElapsedTime(&ms, sc.kernel_start, sc.kernel_stop));
         m_total_kernel_ms += ms;
         m_frames_processed++;
-
-        for (const auto &c : staging)
-            cv.push_back(c);
     }
 };
 

include/aare/clusterfinder_kernel.cuh

@@ -172,19 +172,28 @@ __global__ void find_clusters_in_single_frame(
     if (!valid_pixel)
         return;
 
-    // Per-pixel RMS from global pedestal arrays
-    // rms = sqrt( E[X^2] - E[X]^2 )
+    // Per-pixel variance from global pedestal arrays
+    // Variance = rms^2 = E[X^2] - E[X]^2
+    // NOTE: Keep thresholds squared to avoid one sqrtf() per pixel.
     PEDESTAL_TYPE mean_px = d_pd_mean[global_tid];
     PEDESTAL_TYPE var_px =
         d_pd_sum2[global_tid] / n_pd_samples - mean_px * mean_px;
     PEDESTAL_TYPE rms_sq = max(var_px, PEDESTAL_TYPE{0}); // variance = rms^2
-    COMPUTE_TYPE rms_px = sqrtf(static_cast<COMPUTE_TYPE>(rms_sq));
+    PEDESTAL_TYPE nSig_sq_rms_sq = static_cast<PEDESTAL_TYPE>(m_nSigma) *
+                                   static_cast<PEDESTAL_TYPE>(m_nSigma) *
+                                   rms_sq;
 
     // Pedestal-subtracted value of the center pixel (already in shmem)
     COMPUTE_TYPE val_pixel = shmem[shmem_tid];
 
-    // Negative pedestal early exit
-    if (val_pixel < -m_nSigma * rms_px)
+    // Negative pedestal early exit:
+    //     val_pixel < -nSigma * rms
+    // is equivalent to:
+    //     val_pixel < 0 && val_pixel^2 > nSigma^2 * rms^2
+    if (val_pixel < COMPUTE_TYPE{0} &&
+        static_cast<PEDESTAL_TYPE>(val_pixel) *
+                static_cast<PEDESTAL_TYPE>(val_pixel) >
+            nSig_sq_rms_sq)
         return; // NOTE: pedestal update for this pixel is skipped (same as
                 // sequential)
 
@@ -199,9 +208,6 @@ __global__ void find_clusters_in_single_frame(
     // (ir>=0, ic<=0) PEDESTAL_TYPE br = PEDESTAL_TYPE{0};   // bottom-right
     // (ir>=0, ic>=0)
 
-    CT clusterData[CSX * CSY];
-    int idx = 0; // tracks the pixels in the cluster
-
 #pragma unroll
     for (int ir = -row_radius; ir <= row_radius; ++ir) {
 #pragma unroll
@@ -215,15 +221,6 @@ __global__ void find_clusters_in_single_frame(
             // quadrants) if (ir <= 0 && ic <= 0) tl += val; if (ir <= 0 && ic
             // >= 0) tr += val; if (ir >= 0 && ic <= 0) bl += val; if (ir >= 0
             // && ic >= 0) br += val;
-
-            // Store pedestal-subtracted value in register array for later
-            // cluster output
-            if constexpr (std::is_integral_v<CT>)
-                clusterData[idx] = static_cast<CT>(lroundf(val));
-            else
-                clusterData[idx] = static_cast<CT>(val);
-
-            idx++;
         }
     }
 
@@ -240,12 +237,16 @@ __global__ void find_clusters_in_single_frame(
     //  3. Total significance:         total > c3 * nSigma * rms
     //     -> distributed events where the full cluster sum is significant
 
-    PEDESTAL_TYPE nSig_sq_rms_sq = m_nSigma * m_nSigma * rms_sq;
-
     bool is_photon = false;
 
     // Test 1: single-pixel significance
-    if (max_val > m_nSigma * rms_px) {
+    //     max_val > nSigma * rms
+    // is equivalent to:
+    //     max_val > 0 && max_val^2 > nSigma^2 * rms^2
+    if (max_val > COMPUTE_TYPE{0} &&
+        static_cast<PEDESTAL_TYPE>(max_val) *
+                static_cast<PEDESTAL_TYPE>(max_val) >
+            nSig_sq_rms_sq) {
         // Local-max suppression: only the center-pixel thread records the
         // cluster
         if (val_pixel < max_val)
@@ -293,6 +294,25 @@ __global__ void find_clusters_in_single_frame(
     if (!is_photon) return; // Debugging
     */
 
+    // Delay building clusterData until we know this thread will write a photon.
+    // This avoids CSX*CSY conversions/rounds for the overwhelmingly common
+    // background pixels.
+    CT clusterData[CSX * CSY];
+    int idx = 0;
+
+#pragma unroll
+    for (int ir = -row_radius; ir <= row_radius; ++ir) {
+#pragma unroll
+        for (int ic = -col_radius; ic <= col_radius; ++ic) {
+            COMPUTE_TYPE val = shmem[shmem_tid + ir * shmem_stride + ic];
+            if constexpr (std::is_integral_v<CT>)
+                clusterData[idx] = static_cast<CT>(lroundf(val));
+            else
+                clusterData[idx] = static_cast<CT>(val);
+            idx++;
+        }
+    }
+
     // Write cluster to global output buffer using atomic index
     // for coordination across all blocks
     uint32_t write_idx = atomicAdd(d_cluster_count, 1u);