Files
Jungfraujoch/image_analysis/indexing/FFTIndexerGPU.h
T
leonarski_fandClaude Opus 5.5 3102221568 cuFFT plans kept for the process and lent to one user at a time, output unchanged
Every cuModuleLoadData/Unload in a run is cuFFT making or destroying a plan, and each plan loads
its own kernel modules under a process-wide lock: over a second a plan on a busy GB10, and a PTX
compile per load on Windows. The FFT indexers, rigid-body engines, structure-factor engines,
molecular replacement and beam-centre FFT each made their plans afresh and destroyed them again,
and the sizing functions (rigid-body and structure-factor FFTWorkBytes, beam-centre
DeviceMemoryNeeded via cufftEstimate*) made throwaway plans only to read a work size.

CudaFFTPlan now borrows a plan from a per-process cache keyed by device and full plan shape
(CudaFFTShape: n, embeds, strides, distances, type, batch). A plan is lent to one user at a time
(cuFFT is thread safe only across distinct plans), goes back when its owner is done - after a wait
for the owner's stream - and is never destroyed. Kept plans hold no work area: the user passes its
own, or CudaFFTPlan allocates one from the pool while it is borrowed. The sizing functions read the
work size from the kept plan (CudaFFTPlan::WorkBytes), which the engines then reuse. Plan creation
keeps the NULL-stream wait of 90ba15871, which now also covers the rigid-body, structure-factor
and molecular-replacement plans that were raw cufftHandles without it. A rigid-body engine hands
its previous zone's plans back on SetZone. FitsInDeviceMemory answers false where cuFFT cannot
plan the transform, as the ignored cufftEstimate failure did.

Measured with nsys on an RTX 5080: cuModuleLoadData/Unload 450/450 -> 303/51 on cytc_x10sa,
458/458 -> 161/51 on 4nwv (the remaining unloads are cuFFT's own inside plan creation; the
remaining loads are plans that are alive at the same time, mostly the concurrent rigid-body
engines). Wall unchanged on this workstation (ABAB), where a load costs ~1 ms; the gain is for
GB10 and Windows. Outputs bit-identical (p.hkl and the refined model on 4nwv, cytc_x10sa,
lyso_x06da_ref); cufftPlan3d vs cufftMakePlanMany with a caller work area checked bit-identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-14 21:54:26 +02:00

54 lines
1.6 KiB
C++

// 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;
CudaStream stream;
// After the stream: it waits for the stream before the plan goes back, and goes before the stream.
CudaFFTPlan plan;
void ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) override;
void DirectionsChanged() override;
public:
explicit FFTIndexerGPU(const IndexingSettings& settings);
FFTIndexerGPU(const FFTIndexerGPU &i) = delete;
const FFTIndexerGPU &operator=(const FFTIndexerGPU &i) = delete;
~FFTIndexerGPU() override = default;
};