leonarski_fandClaude Opus 5.5
e8409ef608
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
83 lines
3.0 KiB
C++
83 lines
3.0 KiB
C++
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
|
|
#pragma once
|
|
|
|
#include <vector>
|
|
#include "BraggPrediction.h"
|
|
#include "../indexing/CUDAMemHelpers.h"
|
|
|
|
struct KernelConstsRot {
|
|
float det_width_pxl;
|
|
float det_height_pxl;
|
|
float beam_x;
|
|
float beam_y;
|
|
float coeff_const;
|
|
float one_over_wavelength;
|
|
float one_over_dmax_sq;
|
|
float mos_angle_rad;
|
|
float wedge_angle_rad;
|
|
float min_zeta;
|
|
float mosaicity_multiplier;
|
|
float bandwidth_sigma; // relative dlambda/lambda as a sigma; 0 = monochromatic
|
|
float half_wavelength_A;
|
|
// Angle-dependent sensor efficiency, reduced to the two per-dataset numbers the kernel needs:
|
|
// the optical thickness at normal incidence and the absorbed fraction there. qe_a0 = 0 leaves
|
|
// every reflection untouched (see sensor_absorption::SensorQE).
|
|
float qe_a0;
|
|
float qe_qe0;
|
|
float flight_dL; // normal-incidence air flight path in attenuation lengths; 0 = off
|
|
Coord Astar, Bstar, Cstar, S0;
|
|
Coord m1, m2, m3;
|
|
float m2_S0;
|
|
float m3_S0;
|
|
float rot[9];
|
|
char centering;
|
|
};
|
|
|
|
class BraggPredictionRotGPU : public BraggPrediction {
|
|
// 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;
|
|
|
|
// Device allocations via helpers
|
|
CudaDevicePtr<KernelConstsRot> dK;
|
|
CudaDevicePtr<int> d_count;
|
|
|
|
// 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).
|
|
~BraggPredictionRotGPU() override { cudaStreamSynchronize(stream); }
|
|
|
|
protected:
|
|
void GrowCapacity(int count) override;
|
|
|
|
public:
|
|
int Calc(const DiffractionExperiment &experiment,
|
|
const CrystalLattice &lattice,
|
|
const BraggPredictionSettings &settings) override;
|
|
};
|