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);
}

__global__ void calculate_fft_result(
    const cufftComplex *__restrict__ d_output,
    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;
        size_t offset = static_cast<size_t>(out_len) * 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 += complex_abs(d_output[offset + 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 += complex_abs(d_output[offset + whi]); }
            while (wlo < want_lo) { winsum -= complex_abs(d_output[offset + wlo]); ++wlo; }

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

            const float mag = complex_abs(d_output[offset + 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;
    }
    cudaMemcpy(d_dir_x, dir_x.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_dir_y, dir_y.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_dir_z, dir_z.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice);
}

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

    cudaMemcpyAsync(d_spot_x, spot_x, nspots * sizeof(float), cudaMemcpyHostToDevice, stream);
    cudaMemcpyAsync(d_spot_y, spot_y, nspots * sizeof(float), cudaMemcpyHostToDevice, stream);
    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(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)));

    calculate_fft_result<<<l_gridDim, l_blockDim, 0, stream>>>(d_output_fft,
                                                               max_length_A, min_length_A, histogram_size,
                                                               bg_half,
                                                               direction_vectors.size(), d_result_fft);

    cuda_err(cudaMemcpyAsync(result_fft.data(), d_result_fft, direction_vectors.size() * sizeof(FFTResult),
                             cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));
}