leonarski_fandClaude Opus 5.5
c3cefebe8f
The peak search gave each of the 16384 directions to one thread, a chain of dependent FP64 additions over 1572 bins - latency-bound on a card that holds far more threads. Each direction's spectrum is now split into four consecutive runs of bins, one thread each, and the runs' first-best peaks are combined in order with the same strict comparison, so the first bin of highest prominence is still the one returned. A run after the first sums its starting window afresh rather than carrying the running sum in. The double sums of float magnitudes are exact unless a magnitude is below ~1e-9 of its window sum, so the background is the same double: 0 of 491520 directions differ in a synthetic still/rotation harness, and the hkl of 3000 stills frames (-X fft) is byte-identical. Kernel (isolated, RTX 5080): 0.87 -> 0.65 ms per call. The stills de-novo loop gains ~4% (128 -> 133 frames/s, ABAB, load ~10): the FFT chain as a whole is bound by GPU memory bandwidth under the indexer pool's concurrency, so even a background in float (6x faster kernel, not adopted) gives only ~9%. Rotation first pass: FFT kernels are ~3% of the indexing step. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
329 lines
17 KiB
Plaintext
329 lines
17 KiB
Plaintext
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "FFTIndexerGPU.h"
|
|
#include <cufft.h>
|
|
#include <cmath>
|
|
#include <algorithm>
|
|
|
|
__device__ __host__ inline float complex_abs(const cufftComplex &z) {
|
|
return sqrtf(z.x * z.x + z.y * z.y);
|
|
}
|
|
|
|
// The magnitudes of the transform, transposed to bin-major order: mag_t[bin * directions + direction].
|
|
// The R2C output is direction-major, and the peak search gives each direction to one thread, so read
|
|
// straight from it neighbouring lanes are a whole spectrum apart - every load a separate sector, three
|
|
// streams per thread. Transposed, a warp reads 32 consecutive floats per bin. On an idle card that buys
|
|
// nothing (the search is latency-bound either way, and the transpose costs its own pass); with the
|
|
// stills indexer pool's streams sharing the card it is the difference between 67 and 98 frames/s
|
|
// (measured). A 32x32 tile in shared memory keeps both the read and the write coalesced. The search took
|
|
// complex_abs of the same values, so the magnitudes are bit for bit the ones it computed itself.
|
|
__global__ void magnitude_transpose(const cufftComplex *__restrict__ d_output,
|
|
const int out_len,
|
|
const int directions_size,
|
|
float *__restrict__ mag_t) {
|
|
__shared__ float tile[32][33];
|
|
const int bin0 = blockIdx.x * 32;
|
|
const int dir0 = blockIdx.y * 32;
|
|
for (int r = threadIdx.y; r < 32; r += blockDim.y) {
|
|
const int dir = dir0 + r;
|
|
const int bin = bin0 + threadIdx.x;
|
|
if (dir < directions_size && bin < out_len)
|
|
tile[r][threadIdx.x] = complex_abs(d_output[static_cast<size_t>(dir) * out_len + bin]);
|
|
}
|
|
__syncthreads();
|
|
for (int r = threadIdx.y; r < 32; r += blockDim.y) {
|
|
const int bin = bin0 + r;
|
|
const int dir = dir0 + threadIdx.x;
|
|
if (dir < directions_size && bin < out_len)
|
|
mag_t[static_cast<size_t>(bin) * directions_size + dir] = tile[threadIdx.x][r];
|
|
}
|
|
}
|
|
|
|
// The peak search splits each direction's spectrum into PEAK_CHUNKS consecutive runs of bins and gives each
|
|
// run to its own thread, which keeps a running window sum over its run; one thread per direction leaves the
|
|
// card nearly empty (16384 threads), and this search, a chain of dependent FP64 additions per thread, is
|
|
// then latency-bound. A block holds 32 directions (x, so the bin-major loads stay coalesced) times PEAK_CHUNKS.
|
|
// A run after the first sums its starting window afresh instead of carrying the sum in from the bins before
|
|
// it. That is the same double wherever the sums are exact, and they are: a window of float magnitudes fits
|
|
// a double's 53 bits unless one of them is below ~1e-9 of the window sum (no difference over 3000 stills
|
|
// frames and a rotation-style synthetic set). More runs buy nothing: the card's FP64 rate is then the limit.
|
|
constexpr int PEAK_CHUNKS = 4;
|
|
|
|
__global__ void calculate_fft_result(
|
|
const float *__restrict__ mag_t,
|
|
const float max_length_A,
|
|
const float min_length_A,
|
|
const int histogram_size,
|
|
const int bg_half,
|
|
const int directions_size,
|
|
FFTResult *d_results) {
|
|
__shared__ float run_prom[PEAK_CHUNKS][32];
|
|
__shared__ float run_len[PEAK_CHUNKS][32];
|
|
|
|
const int i = blockIdx.x * 32 + threadIdx.x;
|
|
const int run = threadIdx.y;
|
|
float best_prom = 0.0f;
|
|
float best_len = -1.0f;
|
|
|
|
if (i < directions_size) {
|
|
const int out_len = (histogram_size / 2) + 1;
|
|
// |transform| of bin k of this direction, from the bin-major array (magnitude_transpose)
|
|
const auto mag_at = [&](int k) { return mag_t[static_cast<size_t>(k) * directions_size + i]; };
|
|
const float len_coeff = 2.0f * max_length_A / static_cast<float>(histogram_size);
|
|
|
|
// Pick the peak by PROMINENCE above a running-mean background of half-width bg_half:
|
|
// the projected histogram has a broad low-frequency envelope whose magnitude can
|
|
// exceed the true lattice peaks, so a plain argmax|spec| returns a short envelope
|
|
// vector on weak/pink-beam frames. Subtracting the local mean removes that envelope
|
|
// while sharp lattice peaks keep their height (mirrors FFTIndexerCPU).
|
|
//
|
|
// The window has a constant width, slid inward at the ends instead of truncated (see
|
|
// FFTIndexerCPU): a peak within bg_half of either end - where the LONGEST cells sit - otherwise
|
|
// gets a one-sided background and a biased prominence. A spectrum shorter than the window
|
|
// averages over all of itself.
|
|
const int width = min(2 * bg_half, out_len - 1);
|
|
const int lo_max = out_len - 1 - width;
|
|
const auto window_lo = [&](int j) { return min(max(j - bg_half, 0), lo_max); };
|
|
|
|
// The window width is fixed, so the mean divides by one number throughout. A double division
|
|
// is a long iteration on the FP64 units a consumer card has few of; multiplying by the correctly
|
|
// rounded reciprocal and correcting the quotient once with its exact remainder gives the
|
|
// correctly rounded quotient - the same double the division returns (Markstein (1990) IBM J.
|
|
// Res. Dev. 34, 111-119).
|
|
const double count = width + 1;
|
|
const double inv_count = 1.0 / count;
|
|
|
|
const int j0 = (out_len * run) / PEAK_CHUNKS;
|
|
const int j1 = (out_len * (run + 1)) / PEAK_CHUNKS;
|
|
int lo = window_lo(j0);
|
|
double winsum = 0.0;
|
|
for (int k = lo; k <= lo + width; ++k)
|
|
winsum += mag_at(k);
|
|
|
|
for (int j = j0; j < j1; ++j) {
|
|
if (window_lo(j) > lo) {
|
|
winsum += mag_at(lo + width + 1);
|
|
winsum -= mag_at(lo);
|
|
++lo;
|
|
}
|
|
const float len = len_coeff * static_cast<float>(j);
|
|
if (len <= min_length_A)
|
|
continue;
|
|
const double q = __dmul_rn(winsum, inv_count);
|
|
const float bg = static_cast<float>(__fma_rn(__fma_rn(-q, count, winsum), inv_count, q));
|
|
const float prom = mag_at(j) - bg;
|
|
if (prom > best_prom) {
|
|
best_prom = prom;
|
|
best_len = len;
|
|
}
|
|
}
|
|
}
|
|
|
|
run_prom[run][threadIdx.x] = best_prom;
|
|
run_len[run][threadIdx.x] = best_len;
|
|
__syncthreads();
|
|
|
|
// Each run kept its first bin of highest prominence; taking the runs in order with the same strict
|
|
// comparison gives the direction's first bin of highest prominence.
|
|
if (run == 0 && i < directions_size) {
|
|
FFTResult result{.magnitude = 0.0f, .direction = i, .length = -1};
|
|
for (int r = 0; r < PEAK_CHUNKS; ++r) {
|
|
if (run_prom[r][threadIdx.x] > result.magnitude) {
|
|
result.magnitude = run_prom[r][threadIdx.x];
|
|
result.length = run_len[r][threadIdx.x];
|
|
}
|
|
}
|
|
d_results[i] = result;
|
|
}
|
|
}
|
|
|
|
__global__ void histogram_kernel(const float *__restrict__ coord_x,
|
|
const float *__restrict__ coord_y,
|
|
const float *__restrict__ coord_z,
|
|
const float *__restrict__ dir_x,
|
|
const float *__restrict__ dir_y,
|
|
const float *__restrict__ dir_z,
|
|
float histogram_spacing,
|
|
int histogram_size,
|
|
int coord_size,
|
|
int direction_vectors_size,
|
|
float *__restrict__ output) {
|
|
int direction_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (direction_idx < direction_vectors_size) {
|
|
int base_offset = direction_idx * histogram_size;
|
|
|
|
for (int i = 0; i < histogram_size; i++)
|
|
output[base_offset + i] = 0;
|
|
|
|
for (int i = 0; i < coord_size; i++) {
|
|
float dot = fabsf(
|
|
dir_x[direction_idx] * coord_x[i] + dir_y[direction_idx] * coord_y[i] + dir_z[direction_idx] * coord_z[
|
|
i]);
|
|
int64_t bin = static_cast<int64_t>(dot / histogram_spacing);
|
|
if (bin >= 0 && bin < histogram_size)
|
|
output[base_offset + bin] += 1.0;
|
|
}
|
|
}
|
|
}
|
|
|
|
// The same histogram, one block per direction with the bins in shared memory.
|
|
//
|
|
// The kernel above gives a whole direction to a single thread, so neighbouring lanes write
|
|
// histogram_size floats apart - 12.6 kB at the default sizing. Every warp instruction then touches 32
|
|
// separate sectors of a buffer that is hundreds of megabytes (16384 directions x 3142 bins), with no
|
|
// hope of staying in a 4 MB L2, and the projection loop becomes 160 million scattered global
|
|
// read-modify-writes.
|
|
//
|
|
// The bins are counts, so here they are integers in shared memory. That matters twice: an integer
|
|
// atomicAdd is a real shared-memory instruction on Turing and Ada, where the float one compiles to a
|
|
// compare-and-swap retry loop; and a count below 2^24 converts to float exactly, so the output is bit
|
|
// for bit what the repeated `+= 1.0` above produces. The division by histogram_spacing stays a
|
|
// division - turning it into a multiply by the reciprocal would move a spot across a bin edge.
|
|
__global__ void histogram_shared_kernel(const float *__restrict__ coord_x,
|
|
const float *__restrict__ coord_y,
|
|
const float *__restrict__ coord_z,
|
|
const float *__restrict__ dir_x,
|
|
const float *__restrict__ dir_y,
|
|
const float *__restrict__ dir_z,
|
|
float histogram_spacing,
|
|
int histogram_size,
|
|
int coord_size,
|
|
int direction_vectors_size,
|
|
float *__restrict__ output) {
|
|
extern __shared__ unsigned int bins[];
|
|
const int direction_idx = blockIdx.x;
|
|
if (direction_idx >= direction_vectors_size)
|
|
return;
|
|
|
|
for (int i = threadIdx.x; i < histogram_size; i += blockDim.x)
|
|
bins[i] = 0;
|
|
__syncthreads();
|
|
|
|
const float dx = dir_x[direction_idx], dy = dir_y[direction_idx], dz = dir_z[direction_idx];
|
|
for (int i = threadIdx.x; i < coord_size; i += blockDim.x) {
|
|
const float dot = fabsf(dx * coord_x[i] + dy * coord_y[i] + dz * coord_z[i]);
|
|
const int64_t bin = static_cast<int64_t>(dot / histogram_spacing);
|
|
if (bin >= 0 && bin < histogram_size)
|
|
atomicAdd(&bins[bin], 1u);
|
|
}
|
|
__syncthreads();
|
|
|
|
float *out = output + static_cast<size_t>(direction_idx) * static_cast<size_t>(histogram_size);
|
|
for (int i = threadIdx.x; i < histogram_size; i += blockDim.x)
|
|
out[i] = static_cast<float>(bins[i]);
|
|
}
|
|
|
|
inline void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
|
|
inline void cuda_err(cufftResult val) {
|
|
if (val != cufftResult::CUFFT_SUCCESS)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "CuFFT error");
|
|
}
|
|
|
|
FFTIndexerGPU::FFTIndexerGPU(const IndexingSettings &settings)
|
|
: FFTIndexer(settings), result_fft_reg(result_fft) {
|
|
d_input_fft = CudaDevicePtr<float>(input_size);
|
|
d_output_fft = CudaDevicePtr<cufftComplex>(output_size);
|
|
d_result_fft = CudaDevicePtr<FFTResult>(nDirections);
|
|
|
|
d_spot_x = CudaDevicePtr<float>(FFT_MAX_SPOTS);
|
|
d_spot_y = CudaDevicePtr<float>(FFT_MAX_SPOTS);
|
|
d_spot_z = CudaDevicePtr<float>(FFT_MAX_SPOTS);
|
|
|
|
spot_x = CudaHostPtr<float>(FFT_MAX_SPOTS);
|
|
spot_y = CudaHostPtr<float>(FFT_MAX_SPOTS);
|
|
spot_z = CudaHostPtr<float>(FFT_MAX_SPOTS);
|
|
|
|
d_dir_x = CudaDevicePtr<float>(nDirections);
|
|
d_dir_y = CudaDevicePtr<float>(nDirections);
|
|
d_dir_z = CudaDevicePtr<float>(nDirections);
|
|
|
|
DirectionsChanged();
|
|
|
|
int n[1] = {static_cast<int32_t>(histogram_size)}; // Size of the FFT along a single dimension
|
|
|
|
plan = CudaFFTPlan(1, n, nullptr, 1, histogram_size, nullptr, 1, histogram_size / 2 + 1, CUFFT_R2C,
|
|
nDirections);
|
|
cuda_err(cufftSetStream(plan, stream));
|
|
}
|
|
|
|
|
|
void FFTIndexerGPU::DirectionsChanged() {
|
|
std::vector<float> dir_x(nDirections), dir_y(nDirections), dir_z(nDirections);
|
|
for (int i = 0; i < nDirections; i++) {
|
|
dir_x[i] = direction_vectors.at(i).x;
|
|
dir_y[i] = direction_vectors.at(i).y;
|
|
dir_z[i] = direction_vectors.at(i).z;
|
|
}
|
|
// Checked like every other CUDA call here: a silently failed upload leaves the device holding the
|
|
// PREVIOUS grid, so the directions the host reads the results against are not the ones the kernel
|
|
// used - wrong rows reported as good ones, and nothing says so. On the engine's stream, which the
|
|
// kernels reading them run on: a cudaMemcpy on the NULL stream is not ordered before those.
|
|
cuda_err(cudaMemcpyAsync(d_dir_x, dir_x.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice, stream));
|
|
cuda_err(cudaMemcpyAsync(d_dir_y, dir_y.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice, stream));
|
|
cuda_err(cudaMemcpyAsync(d_dir_z, dir_z.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice, stream));
|
|
}
|
|
|
|
void FFTIndexerGPU::ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) {
|
|
int l_blockDim = 128;
|
|
int l_gridDim = (direction_vectors.size() + l_blockDim - 1) / l_blockDim;
|
|
|
|
for (int i = 0; i < nspots; i++) {
|
|
spot_x[i] = coord[i].x;
|
|
spot_y[i] = coord[i].y;
|
|
spot_z[i] = coord[i].z;
|
|
}
|
|
|
|
cuda_err(cudaMemcpyAsync(d_spot_x, spot_x, nspots * sizeof(float), cudaMemcpyHostToDevice, stream));
|
|
cuda_err(cudaMemcpyAsync(d_spot_y, spot_y, nspots * sizeof(float), cudaMemcpyHostToDevice, stream));
|
|
cuda_err(cudaMemcpyAsync(d_spot_z, spot_z, nspots * sizeof(float), cudaMemcpyHostToDevice, stream));
|
|
|
|
// Shared-memory bins where they fit (they do at any sane sizing - 12.6 kB at the defaults), the
|
|
// thread-per-direction kernel where they do not.
|
|
const size_t hist_shared_bytes = static_cast<size_t>(histogram_size) * sizeof(unsigned int);
|
|
if (hist_shared_bytes <= 48 * 1024) {
|
|
histogram_shared_kernel<<<direction_vectors.size(), 256, hist_shared_bytes, stream>>>(
|
|
d_spot_x, d_spot_y, d_spot_z, d_dir_x, d_dir_y, d_dir_z,
|
|
histogram_spacing, histogram_size, nspots, direction_vectors.size(), d_input_fft);
|
|
} else {
|
|
histogram_kernel<<<l_gridDim, l_blockDim, 0, stream>>>(d_spot_x, d_spot_y, d_spot_z,
|
|
d_dir_x, d_dir_y, d_dir_z,
|
|
histogram_spacing, histogram_size,
|
|
nspots,
|
|
direction_vectors.size(),
|
|
d_input_fft);
|
|
}
|
|
cuda_err(cudaGetLastError());
|
|
|
|
cuda_err(cufftExecR2C(plan, d_input_fft, d_output_fft));
|
|
|
|
// Background half-window ~15 A (length-based, so independent of histogram sizing); see
|
|
// FFTIndexerCPU for the prominence-vs-envelope rationale and the validated optimum.
|
|
const double len_coeff = 2.0 * static_cast<double>(max_length_A) / static_cast<double>(histogram_size);
|
|
const int bg_half = std::max(1, static_cast<int>(std::lround(15.0 / len_coeff)));
|
|
|
|
const int out_len = static_cast<int>(histogram_size / 2 + 1);
|
|
const dim3 tr_grid((out_len + 31) / 32, (direction_vectors.size() + 31) / 32);
|
|
// The histogram has been transformed and is not read again, so its buffer - histogram_size floats
|
|
// per direction, more than the out_len magnitudes need - holds them.
|
|
float *d_mag_t = d_input_fft;
|
|
magnitude_transpose<<<tr_grid, dim3(32, 8), 0, stream>>>(d_output_fft, out_len,
|
|
direction_vectors.size(), d_mag_t);
|
|
cuda_err(cudaGetLastError());
|
|
const int peak_grid = (static_cast<int>(direction_vectors.size()) + 31) / 32;
|
|
calculate_fft_result<<<peak_grid, dim3(32, PEAK_CHUNKS), 0, stream>>>(d_mag_t,
|
|
max_length_A, min_length_A, histogram_size,
|
|
bg_half,
|
|
direction_vectors.size(), d_result_fft);
|
|
cuda_err(cudaGetLastError());
|
|
|
|
cuda_err(cudaMemcpyAsync(result_fft.data(), d_result_fft, direction_vectors.size() * sizeof(FFTResult),
|
|
cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaStreamSynchronize(stream));
|
|
}
|