From e8409ef608c25ec5b9d19702b80ec3fbc6e5e469 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 12 Oct 2026 22:36:31 +0200 Subject: [PATCH] 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) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- .../bragg_prediction/BraggPredictionRotGPU.cu | 144 ++++++++++++++++-- .../bragg_prediction/BraggPredictionRotGPU.h | 21 ++- tests/CalcBraggPredictionTest.cpp | 47 ++++++ 3 files changed, 194 insertions(+), 18 deletions(-) diff --git a/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu b/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu index 83dc1717d..31353f61f 100644 --- a/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu +++ b/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu @@ -10,6 +10,7 @@ #ifdef JFJOCH_USE_CUDA #include "../indexing/CUDAMemHelpers.h" #include +#include #include 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(x.h + p.max_h) << p.shift_h) + | (static_cast(x.k + p.max_k) << p.shift_k) + | static_cast(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(~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(count); - d_out = CudaDevicePtr(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(capacity); + for (int i = 0; i < 2; i++) { + d_key32[i] = CudaDevicePtr(capacity); + d_key64[i] = CudaDevicePtr(capacity); + d_index[i] = CudaDevicePtr(capacity); + } + size_t bytes32 = 0, bytes64 = 0; + cub::DoubleBuffer k32(d_key32[0], d_key32[1]); + cub::DoubleBuffer k64(d_key64[0], d_key64[1]); + cub::DoubleBuffer 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(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(kept); + h_out = CudaHostPtr(kept); + } + if (reflections.size() < static_cast(kept)) + reflections.resize(kept); + + const int threads = 256; + const int blocks = (count + threads - 1) / threads; + cub::DoubleBuffer k32(d_key32[0], d_key32[1]); + cub::DoubleBuffer k64(d_key64[0], d_key64[1]); + cub::DoubleBuffer 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<<>>(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<<>>(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<<>>(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 \ No newline at end of file diff --git a/image_analysis/bragg_prediction/BraggPredictionRotGPU.h b/image_analysis/bragg_prediction/BraggPredictionRotGPU.h index d56dd9c86..ffafcadea 100644 --- a/image_analysis/bragg_prediction/BraggPredictionRotGPU.h +++ b/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 h_out; + // The device's reflections, in the order the atomics placed them. CudaDevicePtr 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 d_kept; + CudaHostPtr h_out; + int kept_capacity = 0; + + // Sort keys and the permutation they carry, double-buffered for cub's radix sort, and its scratch. + CudaDevicePtr d_key32[2]; + CudaDevicePtr d_key64[2]; + CudaDevicePtr d_index[2]; + CudaDevicePtr 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 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). diff --git a/tests/CalcBraggPredictionTest.cpp b/tests/CalcBraggPredictionTest.cpp index 1daec3df4..465a6d4d7 100644 --- a/tests/CalcBraggPredictionTest.cpp +++ b/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 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 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