BraggPredictionRotGPU: order and select the predictions on the device

The rotation predictor used to download every predicted reflection into pinned memory, copy it into
the host vector, build sort keys, std::sort them, gather the 88-byte structs into scratch and copy
them back (OrderOutput), and - on a frame over output_limit - partial_sort the structs themselves
(TruncateToOutput). On a large cell that is a quarter to a third of the image-loop worker CPU
(profiles of 4nwv and 7qij: TruncateToOutput 12-16 %, OrderOutput 9-13 %, plus the memmoves).

Now the same order is built on the device with stable cub radix sorts of an index: delta_phi, then
the packed hkl (OrderOutput's h, k, l, delta_phi), then - over the cap only - partiality descending
and excitation error ascending (TruncateToOutput, whose hkl tie-break the stable sort keeps). Only
the kept rows are gathered and downloaded. Floats are keyed by an order-preserving integer map with
-0 made +0, so the key order is the comparators' order; the only possible difference from the host
is between two solutions of the SAME hkl with bit-identical partiality and excitation error, whose
relative order partial_sort leaves to its heap and the stable sort to delta_phi.

The device buffers grow with half again of slack, so a frame predicting a few more than the last
record (the supercell probe's frames above all) no longer reallocates and predicts again; the host
side is sized for what is returned, at most output_limit.

Test: BraggPredictionRotGPU_DeviceOrderMatchesHost holds the device output against OrderOutput's
order and against TruncateToOutput applied on the host, at four caps.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3

image_analysis/bragg_prediction/BraggPredictionRotGPU.cu

@@ -10,6 +10,7 @@
 #ifdef JFJOCH_USE_CUDA
 #include "../indexing/CUDAMemHelpers.h"
 #include <cuda_runtime.h>
+#include <cub/cub.cuh>
 #include <cmath>
 
 static inline void cuda_err(cudaError_t val) {
@@ -256,6 +257,60 @@ namespace {
         }
     }
 
+    // A float as an unsigned integer that sorts in the same order. -0 is first made +0, which the
+    // comparators of OrderOutput and TruncateToOutput treat as equal to it.
+    __device__ inline uint32_t FloatOrderKey(float v) {
+        const uint32_t u = __float_as_uint(v + 0.0f);
+        return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
+    }
+
+    struct HklPacking {
+        int max_h, max_k, max_l;
+        int shift_h, shift_k;   // l takes the lowest bits
+    };
+
+    __global__ void phi_key_kernel(const Reflection *__restrict__ r, int n,
+                                   uint32_t *__restrict__ key, int *__restrict__ index) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i >= n) return;
+        key[i] = FloatOrderKey(r[i].delta_phi_deg);
+        index[i] = i;
+    }
+
+    __global__ void hkl_key_kernel(const Reflection *__restrict__ r, const int *__restrict__ index, int n,
+                                   HklPacking p, uint64_t *__restrict__ key) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i >= n) return;
+        const Reflection &x = r[index[i]];
+        key[i] = (static_cast<uint64_t>(x.h + p.max_h) << p.shift_h)
+                 | (static_cast<uint64_t>(x.k + p.max_k) << p.shift_k)
+                 | static_cast<uint64_t>(x.l + p.max_l);
+    }
+
+    // Largest partiality first, then smallest excitation error (see TruncateToOutput).
+    __global__ void best_recorded_key_kernel(const Reflection *__restrict__ r, const int *__restrict__ index,
+                                             int n, uint64_t *__restrict__ key) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i >= n) return;
+        const Reflection &x = r[index[i]];
+        key[i] = (static_cast<uint64_t>(~FloatOrderKey(x.partiality)) << 32) | FloatOrderKey(x.dist_ewald);
+    }
+
+    __global__ void gather_kernel(const Reflection *__restrict__ r, const int *__restrict__ index, int n,
+                                  Reflection *__restrict__ out) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i >= n) return;
+        out[i] = r[index[i]];
+    }
+
+    // Bits that hold 0 .. 2 * max_index.
+    int IndexBits(int max_index) {
+        int bits = 1;
+        while ((1ll << bits) < 2ll * max_index + 1)
+            bits++;
+        return bits;
+    }
+
     inline KernelConstsRot BuildKernelConstsRot(const DiffractionExperiment &experiment,
                                                 const CrystalLattice &lattice,
                                                 const BraggPredictionSettings &settings) {
@@ -333,16 +388,81 @@ namespace {
 } // namespace
 
 BraggPredictionRotGPU::BraggPredictionRotGPU(int max_reflections)
-    : BraggPrediction(max_reflections),
-      h_out(max_reflections), d_out(max_reflections),
-      dK(1), d_count(1), h_count(1) {
+    : BraggPrediction(0), dK(1), d_count(1), h_count(1) {
+    GrowCapacity(max_reflections);
 }
 
-// The host buffer is page-locked and the device buffer is sized to match it, so both are rebuilt.
+// Only the device side follows the prediction count: the host gets the kept rows alone, so
+// `reflections` is sized in OrderAndSelect. Grown with slack, so a frame predicting a few more than
+// the last record - the supercell probe's frames above all - does not reallocate every buffer and
+// predict again.
 void BraggPredictionRotGPU::GrowCapacity(int count) {
-    BraggPrediction::GrowCapacity(count);
-    h_out = CudaHostPtr<Reflection>(count);
-    d_out = CudaDevicePtr<Reflection>(count);
+    const int capacity = count + count / 2;
+    max_reflections = capacity;
+    // Queued work may still use the buffers about to be freed (see CudaDevicePtr).
+    cudaStreamSynchronize(stream);
+    d_out = CudaDevicePtr<Reflection>(capacity);
+    for (int i = 0; i < 2; i++) {
+        d_key32[i] = CudaDevicePtr<uint32_t>(capacity);
+        d_key64[i] = CudaDevicePtr<uint64_t>(capacity);
+        d_index[i] = CudaDevicePtr<int>(capacity);
+    }
+    size_t bytes32 = 0, bytes64 = 0;
+    cub::DoubleBuffer<uint32_t> k32(d_key32[0], d_key32[1]);
+    cub::DoubleBuffer<uint64_t> k64(d_key64[0], d_key64[1]);
+    cub::DoubleBuffer<int> idx(d_index[0], d_index[1]);
+    cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes32, k32, idx, capacity, 0, 32, stream));
+    cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes64, k64, idx, capacity, 0, 64, stream));
+    sort_temp_bytes = std::max(bytes32, bytes64);
+    d_sort_temp = CudaDevicePtr<uint8_t>(sort_temp_bytes);
+}
+
+int BraggPredictionRotGPU::OrderAndSelect(int count, const BraggPredictionSettings &settings) {
+    const int kept = std::min(count, output_limit);
+    if (kept > kept_capacity) {
+        cudaStreamSynchronize(stream);
+        kept_capacity = kept;
+        d_kept = CudaDevicePtr<Reflection>(kept);
+        h_out = CudaHostPtr<Reflection>(kept);
+    }
+    if (reflections.size() < static_cast<size_t>(kept))
+        reflections.resize(kept);
+
+    const int threads = 256;
+    const int blocks = (count + threads - 1) / threads;
+    cub::DoubleBuffer<uint32_t> k32(d_key32[0], d_key32[1]);
+    cub::DoubleBuffer<uint64_t> k64(d_key64[0], d_key64[1]);
+    cub::DoubleBuffer<int> idx(d_index[0], d_index[1]);
+    size_t temp_bytes = sort_temp_bytes;
+
+    // OrderOutput's order - h, k, l, then delta_phi - as two stable radix sorts, the minor key first.
+    phi_key_kernel<<<blocks, threads, 0, stream>>>(d_out, count, k32.Current(), idx.Current());
+    cuda_err(cudaGetLastError());
+    cuda_err(cub::DeviceRadixSort::SortPairs(d_sort_temp, temp_bytes, k32, idx, count, 0, 32, stream));
+
+    HklPacking pack{settings.max_h, settings.max_k, settings.max_l};
+    const int bits_l = IndexBits(settings.max_l);
+    const int bits_k = IndexBits(settings.max_k);
+    pack.shift_k = bits_l;
+    pack.shift_h = bits_l + bits_k;
+    const int hkl_bits = pack.shift_h + IndexBits(settings.max_h);
+    hkl_key_kernel<<<blocks, threads, 0, stream>>>(d_out, idx.Current(), count, pack, k64.Current());
+    cuda_err(cudaGetLastError());
+    cuda_err(cub::DeviceRadixSort::SortPairs(d_sort_temp, temp_bytes, k64, idx, count, 0, hkl_bits, stream));
+
+    // Over the cap, TruncateToOutput's order on top: a stable sort leaves hkl breaking its ties.
+    if (count > output_limit) {
+        best_recorded_key_kernel<<<blocks, threads, 0, stream>>>(d_out, idx.Current(), count, k64.Current());
+        cuda_err(cudaGetLastError());
+        cuda_err(cub::DeviceRadixSort::SortPairs(d_sort_temp, temp_bytes, k64, idx, count, 0, 64, stream));
+    }
+
+    gather_kernel<<<(kept + threads - 1) / threads, threads, 0, stream>>>(d_out, idx.Current(), kept, d_kept);
+    cuda_err(cudaGetLastError());
+    cuda_err(cudaMemcpyAsync(h_out, d_kept, sizeof(Reflection) * kept, cudaMemcpyDeviceToHost, stream));
+    cuda_err(cudaStreamSynchronize(stream));
+    std::copy(h_out.get(), h_out.get() + kept, reflections.begin());
+    return kept;
 }
 
 int BraggPredictionRotGPU::Calc(const DiffractionExperiment &experiment,
@@ -380,13 +500,7 @@ int BraggPredictionRotGPU::Calc(const DiffractionExperiment &experiment,
         count = std::min(*h_count.get(), max_reflections);
     }
     if (count == 0) return 0;
-
-    cudaMemcpyAsync(h_out, d_out, sizeof(Reflection) * count, cudaMemcpyDeviceToHost, stream);
-    cudaStreamSynchronize(stream);
-    std::copy(h_out.get(), h_out.get() + count, reflections.begin());
-
-    OrderOutput(count);
-    return TruncateToOutput(count);
+    return OrderAndSelect(count, settings);
 }
 
-#endif
+#endif

image_analysis/bragg_prediction/BraggPredictionRotGPU.h

@@ -37,10 +37,21 @@ struct KernelConstsRot {
 };
 
 class BraggPredictionRotGPU : public BraggPrediction {
-    // The device's reflections land here and are copied into `reflections`: pinned, so the download is a
-    // DMA, and from the pool, so the engine does not register and unregister a vector of its own.
-    CudaHostPtr<Reflection> h_out;
+    // The device's reflections, in the order the atomics placed them.
     CudaDevicePtr<Reflection> d_out;
+    // The rows Calc returns, ordered and selected on the device, then downloaded into h_out: pinned, so
+    // the download is a DMA, and from the pool, so the engine does not register and unregister a vector
+    // of its own. Sized for what is returned (at most output_limit), not for what was predicted.
+    CudaDevicePtr<Reflection> d_kept;
+    CudaHostPtr<Reflection> h_out;
+    int kept_capacity = 0;
+
+    // Sort keys and the permutation they carry, double-buffered for cub's radix sort, and its scratch.
+    CudaDevicePtr<uint32_t> d_key32[2];
+    CudaDevicePtr<uint64_t> d_key64[2];
+    CudaDevicePtr<int> d_index[2];
+    CudaDevicePtr<uint8_t> d_sort_temp;
+    size_t sort_temp_bytes = 0;
 
     // Dedicated stream
     CudaStream stream;
@@ -52,6 +63,10 @@ class BraggPredictionRotGPU : public BraggPrediction {
     // Host pinned buffer for async download (optional but faster)
     CudaHostPtr<int> h_count;
 
+    // Order the `count` device reflections exactly as OrderOutput and TruncateToOutput would, keep the
+    // first min(count, output_limit) and download them into `reflections`. Returns the kept count.
+    int OrderAndSelect(int count, const BraggPredictionSettings &settings);
+
 public:
     explicit BraggPredictionRotGPU(int max_reflections = kPredictionCapacity);
     // Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).

tests/CalcBraggPredictionTest.cpp

@@ -861,4 +861,51 @@ TEST_CASE("BraggPrediction_CPU_GPU_consistency_large_cell") {
         CHECK(matched >= gpu_count - gpu_count / 1000);
     }
 }
+
+TEST_CASE("BraggPredictionRotGPU_DeviceOrderMatchesHost") {
+    // The GPU predictor orders and selects its reflections on the device. What it returns has to be what
+    // OrderOutput and TruncateToOutput make of the same predictions on the host: the full list in hkl,
+    // delta_phi order, and under a cap the best-recorded ones in TruncateToOutput's order.
+    const auto experiment = LargeCellExperiment();
+    const auto lattice = LargeCell();
+    const auto settings = LargeCellSettings();
+
+    BraggPredictionRotGPU gpu;
+    gpu.output_limit = 1 << 30;
+    const int n = gpu.Calc(experiment, lattice, settings);
+    REQUIRE(n > BraggPrediction::kPredictionCapacity);
+    const std::vector<Reflection> all(gpu.GetReflections().begin(), gpu.GetReflections().begin() + n);
+    CHECK(std::is_sorted(all.begin(), all.end(), [](const Reflection &a, const Reflection &b) {
+        if (a.h != b.h) return a.h < b.h;
+        if (a.k != b.k) return a.k < b.k;
+        if (a.l != b.l) return a.l < b.l;
+        return a.delta_phi_deg < b.delta_phi_deg;
+    }));
+
+    for (const int limit : {n - 1, n / 2, n / 7, 1}) {
+        std::vector<Reflection> host = all;
+        std::partial_sort(host.begin(), host.begin() + limit, host.end(),
+                          [](const Reflection &a, const Reflection &b) {
+                              if (a.partiality != b.partiality) return a.partiality > b.partiality;
+                              if (a.dist_ewald != b.dist_ewald) return a.dist_ewald < b.dist_ewald;
+                              if (a.h != b.h) return a.h < b.h;
+                              if (a.k != b.k) return a.k < b.k;
+                              return a.l < b.l;
+                          });
+        gpu.output_limit = limit;
+        REQUIRE(gpu.Calc(experiment, lattice, settings) == limit);
+        INFO("limit " << limit);
+        // Field by field: the kernel leaves the intensity fields unset, so the bytes of two runs differ there.
+        auto same = [](const Reflection &a, const Reflection &b) {
+            return a.h == b.h && a.k == b.k && a.l == b.l && a.delta_phi_deg == b.delta_phi_deg
+                   && a.predicted_x == b.predicted_x && a.predicted_y == b.predicted_y && a.d == b.d
+                   && a.dist_ewald == b.dist_ewald && a.partiality == b.partiality && a.zeta == b.zeta
+                   && a.image_scale_corr == b.image_scale_corr;
+        };
+        size_t ndiff = 0;
+        for (int i = 0; i < limit; i++)
+            ndiff += !same(host[i], gpu.GetReflections()[i]);
+        CHECK(ndiff == 0);
+    }
+}
 #endif