Jungfraujoch/image_analysis/indexing/FFTIndexerGPU.h

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

// This include should be only included in sections of the code, where it is certain that CUDA is present
// so with JFJOCH_USE_CUDA preprocessor definition, given this file is included in the source only in this case

#include <vector>
#include <mutex>
#include <optional>
#include "CUDAMemHelpers.h"
#include "../../common/Coord.h"
#include "../../common/CrystalLattice.h"
#include "FFTIndexer.h"
#include "../common/IndexingSettings.h"
#include "FFTResult.h"

class FFTIndexerGPU : public FFTIndexer {
    CudaDevicePtr<float> d_dir_x;
    CudaDevicePtr<float> d_dir_y;
    CudaDevicePtr<float> d_dir_z;

    CudaDevicePtr<float> d_spot_x;
    CudaDevicePtr<float> d_spot_y;
    CudaDevicePtr<float> d_spot_z;

    CudaHostPtr<float> spot_x;
    CudaHostPtr<float> spot_y;
    CudaHostPtr<float> spot_z;

    CudaDevicePtr<float> d_input_fft;
    CudaDevicePtr<cufftComplex> d_output_fft;
    CudaDevicePtr<FFTResult> d_result_fft;

    CudaRegisteredVector<FFTResult> result_fft_reg;

    CudaFFTPlan plan;

    CudaStream stream;

    void ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) override;
public:
    explicit FFTIndexerGPU(const IndexingSettings& settings);
    FFTIndexerGPU(const FFTIndexerGPU &i) = delete;
    const FFTIndexerGPU &operator=(const FFTIndexerGPU &i) = delete;
    ~FFTIndexerGPU() override = default;
};