Jungfraujoch/image_analysis/indexing/FFTIndexerGPU.cu

// 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];
    }
}

__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) {
    int i = blockIdx.x * blockDim.x + threadIdx.x; // Get thread index

    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]; };
        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).
        double winsum = 0.0;
        int wlo = 0;
        int whi = min(bg_half, out_len - 1);
        for (int k = wlo; k <= whi; ++k)
            winsum += mag_at(k);

        float best_prom = 0.0f;
        FFTResult result{.magnitude = 0.0f, .direction = i, .length = -1};

        for (int j = 0; j < out_len; ++j) {
            // Constant-width window, 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. Both bounds stay
            // monotonically non-decreasing in j, so the running sum below is still valid.
            int want_lo = j - bg_half;
            int want_hi = j + bg_half;
            if (want_lo < 0) { want_hi = min(out_len - 1, want_hi - want_lo); want_lo = 0; }
            if (want_hi > out_len - 1) {
                want_lo = max(0, want_lo - (want_hi - (out_len - 1)));
                want_hi = out_len - 1;
            }
            while (whi < want_hi) { ++whi; winsum += mag_at(whi); }
            while (wlo < want_lo) { winsum -= mag_at(wlo); ++wlo; }

            const float len = len_coeff * static_cast<float>(j);
            if (len <= min_length_A) continue;

            const float mag = mag_at(j);
            const float bg = static_cast<float>(winsum / static_cast<double>(whi - wlo + 1));
            const float prom = mag - bg;
            if (prom > best_prom) {
                best_prom = prom;
                result.magnitude = prom;
                result.length = len;
            }
        }
        d_results[i] = result; // Store the 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());
    calculate_fft_result<<<l_gridDim, l_blockDim, 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));
}