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>

image_analysis/geom_refinement/BeamCenterFFTGPU.cu

@@ -45,6 +45,33 @@ void Check(cufftResult res, const char *what) {
 std::vector<int> Dims(int64_t a) { return {static_cast<int>(a)}; }
 std::vector<int> Dims(int64_t a, int64_t b) { return {static_cast<int>(a), static_cast<int>(b)}; }
 
+// The transforms, as both the computation and its memory estimate take them: the point surfaces' 2D
+// pair on an ny x nx grid, and the line surfaces' 1D pair of length nfft.
+CudaFFTShape PointForward(int64_t ny, int64_t nx) {
+    const int64_t nxc = nx / 2 + 1;
+    return {.n = Dims(ny, nx), .inembed = Dims(ny, nx), .idist = static_cast<int>(ny * nx),
+            .onembed = Dims(ny, nxc), .odist = static_cast<int>(ny * nxc), .type = CUFFT_R2C};
+}
+
+CudaFFTShape PointInverse(int64_t ny, int64_t nx) {
+    const int64_t nxc = nx / 2 + 1;
+    return {.n = Dims(ny, nx), .inembed = Dims(ny, nxc), .idist = static_cast<int>(ny * nxc),
+            .onembed = Dims(ny, nx), .odist = static_cast<int>(ny * nx), .type = CUFFT_C2R};
+}
+
+CudaFFTShape LineForward(int64_t nfft, int64_t nbatch) {
+    const int64_t nc = nfft / 2 + 1;
+    return {.n = Dims(nfft), .inembed = Dims(nfft), .idist = static_cast<int>(nfft),
+            .onembed = Dims(nc), .odist = static_cast<int>(nc), .type = CUFFT_R2C,
+            .batch = static_cast<int>(nbatch)};
+}
+
+CudaFFTShape LineInverse(int64_t nfft) {
+    const int64_t nc = nfft / 2 + 1;
+    return {.n = Dims(nfft), .inembed = Dims(nc), .idist = static_cast<int>(nc),
+            .onembed = Dims(nfft), .odist = static_cast<int>(nfft), .type = CUFFT_C2R, .batch = 4};
+}
+
 __global__ void SquareInPlace(float *data, int64_t n) {
     const int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
     if (i < n)
@@ -249,10 +276,8 @@ PointSurfacesDevice TransformPoint(const std::vector<float> &a, const std::vecto
     out.real = CudaDevicePtr<float>(nreal, ALLOC);
     CudaDevicePtr<cufftComplex> d_a(ncomplex, ALLOC), d_m(ncomplex, ALLOC), d_prod(ncomplex, ALLOC);
 
-    CudaFFTPlan fwd(2, Dims(ny, nx), Dims(ny, nx), 1, static_cast<int>(nreal), Dims(ny, nxc), 1,
-                    static_cast<int>(ncomplex), CUFFT_R2C, 1);
-    CudaFFTPlan inv(2, Dims(ny, nx), Dims(ny, nxc), 1, static_cast<int>(ncomplex), Dims(ny, nx), 1,
-                    static_cast<int>(nreal), CUFFT_C2R, 1);
+    CudaFFTPlan fwd(PointForward(ny, nx));
+    CudaFFTPlan inv(PointInverse(ny, nx));
 
     auto upload = [&](const std::vector<float> &src) {
         Check(cudaMemset(out.real.get(), 0, nreal * sizeof(float)),
@@ -428,8 +453,7 @@ BeamCenterConvSurfaces1D BeamCenterFFTGPU::LineSurfaces(const std::vector<float>
     {
         CudaDevicePtr<float> d_img(h * w, ALLOC), d_in(nfft * nbatch, ALLOC);
         CudaDevicePtr<cufftComplex> d_a(nbatch * nc, ALLOC), d_m(nbatch * nc, ALLOC), d_a2(nbatch * nc, ALLOC);
-        CudaFFTPlan fwd(1, Dims(nfft), Dims(nfft), 1, static_cast<int>(nfft), Dims(nc), 1,
-                        static_cast<int>(nc), CUFFT_R2C, static_cast<int>(nbatch));
+        CudaFFTPlan fwd(LineForward(nfft, nbatch));
 
         auto transform = [&](int square, cufftComplex *dst) {
             GatherSequences<<<Blocks(nfft * nbatch), BLOCK>>>(d_img.get(), nseq, nbatch, stride_seq,
@@ -459,8 +483,7 @@ BeamCenterConvSurfaces1D BeamCenterFFTGPU::LineSurfaces(const std::vector<float>
     CudaDevicePtr<float> d_out(4 * nfft, ALLOC);
     PackAccumulators<<<Blocks(nc), BLOCK>>>(accC, accS, accQ, accD, nc, d_pack.get());
     CheckLastKernel("BeamCenterFFT: accumulator packing failed");
-    CudaFFTPlan inv(1, Dims(nfft), Dims(nc), 1, static_cast<int>(nc), Dims(nfft), 1,
-                    static_cast<int>(nfft), CUFFT_C2R, 4);
+    CudaFFTPlan inv(LineInverse(nfft));
     Check(cufftExecC2R(inv, d_pack, d_out), "BeamCenterFFT: inverse transform failed");
 
     std::vector<float> host(static_cast<size_t>(4 * nfft));
@@ -486,9 +509,8 @@ size_t BeamCenterFFTGPU::DeviceMemoryNeeded(int64_t width, int64_t height) {
     const int64_t ny = BeamCenterFFTPadSize(2 * height);
     const int64_t nx = BeamCenterFFTPadSize(2 * width);
     const int64_t nxc = nx / 2 + 1;
-    size_t work_fwd = 0, work_inv = 0;
-    cufftEstimate2d(static_cast<int>(ny), static_cast<int>(nx), CUFFT_R2C, &work_fwd);
-    cufftEstimate2d(static_cast<int>(ny), static_cast<int>(nx), CUFFT_C2R, &work_inv);
+    const size_t work_fwd = CudaFFTPlan::WorkBytes(PointForward(ny, nx));
+    const size_t work_inv = CudaFFTPlan::WorkBytes(PointInverse(ny, nx));
     // One transform buffer, three spectra, and three of the four convolution surfaces - the
     // fourth stays in the transform buffer, where the last inverse leaves it.
     size_t needed = ny * nx * sizeof(float) + 3 * ny * nxc * sizeof(cufftComplex)
@@ -502,11 +524,7 @@ size_t BeamCenterFFTGPU::DeviceMemoryNeeded(int64_t width, int64_t height) {
         const int64_t nbatch = pass == 0 ? width : height;
         const int64_t nfft = BeamCenterFFTPadSize(2 * nseq);
         const int64_t nc = nfft / 2 + 1;
-        std::vector<int> n = Dims(nfft), embed_real = Dims(nfft), embed_complex = Dims(nc);
-        size_t work = 0;
-        cufftEstimateMany(1, n.data(), embed_real.data(), 1, static_cast<int>(nfft),
-                          embed_complex.data(), 1, static_cast<int>(nc), CUFFT_R2C,
-                          static_cast<int>(nbatch), &work);
+        const size_t work = CudaFFTPlan::WorkBytes(LineForward(nfft, nbatch));
         const size_t line = (height * width + nfft * nbatch) * sizeof(float)
                             + 3 * nbatch * nc * sizeof(cufftComplex) + work;
         needed = std::max(needed, line);
@@ -515,11 +533,19 @@ size_t BeamCenterFFTGPU::DeviceMemoryNeeded(int64_t width, int64_t height) {
 }
 
 bool BeamCenterFFTGPU::FitsInDeviceMemory(int64_t width, int64_t height) {
+    // The work areas are read from the plans themselves, and a transform cuFFT cannot plan does not fit.
+    size_t needed = 0;
+    try {
+        needed = DeviceMemoryNeeded(width, height);
+    } catch (const JFJochException &) {
+        cuda_clear_error();
+        return false;
+    }
     size_t free_bytes = 0, total_bytes = 0;
     if (cudaMemGetInfo(&free_bytes, &total_bytes) != cudaSuccess)
         return false;
     // Headroom for the allocator's own rounding, and for whatever else lands on a shared card
     // between this question and the allocations it is asked about.
     constexpr size_t HEADROOM = 512ull << 20;
-    return DeviceMemoryNeeded(width, height) + HEADROOM <= free_bytes;
+    return needed + HEADROOM <= free_bytes;
 }

image_analysis/indexing/CUDAMemHelpers.h

@@ -5,6 +5,7 @@
 
 #include <cuda_runtime.h>
 #include <cufft.h>
+#include <algorithm>
 #include <cstdio>
 #include <map>
 #include <mutex>
@@ -106,67 +107,6 @@ public:
     cudaEvent_t get() const { return event_; }
 };
 
-class CudaFFTPlan {
-    cufftHandle plan_ = 0;
-public:
-    CudaFFTPlan() = default;
-
-    CudaFFTPlan(
-            int rank,
-            const int* n,
-            const int* inembed, int istride, int idist,
-            const int* onembed, int ostride, int odist,
-            cufftType type, int batch)
-    {
-        const cufftResult r = cufftPlanMany(&plan_, rank, const_cast<int*>(n),
-                                            const_cast<int*>(inembed), istride, idist,
-                                            const_cast<int*>(onembed), ostride, odist,
-                                            type, batch);
-        // Said as a memory failure whatever the code: on a full card (two runs sharing it) the plan
-        // failed as CUFFT_INTERNAL_ERROR, not CUFFT_ALLOC_FAILED.
-        if (r != CUFFT_SUCCESS)
-            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
-                                  GPUMemoryFailure("Creating the CUFFT plan (cufftPlanMany)", 0, true,
-                                                   ("cuFFT error " + std::to_string(r)).c_str()));
-        // The plan's setup can still be queued on the NULL stream when cufftPlanMany returns, and a
-        // non-blocking stream (CudaStream) does not wait for it: a transform launched on one straight
-        // away read the plan's tables before they were there - an illegal address in 3 of 4 runs where
-        // a rotation pool built its FFT indexers at transform length 3072. Waiting here orders it.
-        if (cudaStreamSynchronize(nullptr) != cudaSuccess)
-            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
-                                  "Creating the CUFFT plan: setup on the NULL stream failed");
-    }
-
-    // Convenience overload for vector input
-    CudaFFTPlan(
-            int rank,
-            const std::vector<int>& n,
-            const std::vector<int>& inembed, int istride, int idist,
-            const std::vector<int>& onembed, int ostride, int odist,
-            cufftType type, int batch)
-            : CudaFFTPlan(rank, n.data(), inembed.data(), istride, idist, onembed.data(), ostride, odist, type, batch)
-    {}
-
-    ~CudaFFTPlan() {
-        if (plan_) cufftDestroy(plan_);
-    }
-    // Move-only type
-    CudaFFTPlan(CudaFFTPlan&& other) noexcept : plan_(other.plan_) { other.plan_ = 0; }
-    CudaFFTPlan& operator=(CudaFFTPlan&& other) noexcept {
-        if (this != &other) {
-            if (plan_) cufftDestroy(plan_);
-            plan_ = other.plan_;
-            other.plan_ = 0;
-        }
-        return *this;
-    }
-    CudaFFTPlan(const CudaFFTPlan&) = delete;
-    CudaFFTPlan& operator=(const CudaFFTPlan&) = delete;
-
-    operator cufftHandle() const { return plan_; }
-    cufftHandle get() const { return plan_; }
-};
-
 // How much freed memory the pool keeps rather than returning it to the driver. Returning it puts the
 // next allocation straight back on the path this exists to avoid, so hold the per-worker engine
 // buffers; but the card also has to fit the merge, which asks for several gigabytes of its own after
@@ -389,6 +329,180 @@ public:
     operator T*() const { return ptr_; }
 };
 
+// The transform a cuFFT plan is made for, as cufftMakePlanMany takes it; an empty embed is its nullptr.
+struct CudaFFTShape {
+    std::vector<int> n;
+    std::vector<int> inembed;
+    int istride = 1, idist = 0;
+    std::vector<int> onembed;
+    int ostride = 1, odist = 0;
+    cufftType type = CUFFT_R2C;
+    int batch = 1;
+};
+
+// cuFFT plans are kept for the life of the process, by device and shape, and lent to one user at a time.
+// Making a plan is costly and serialised: cuFFT holds a process-wide lock while it loads the plan's
+// kernel modules (each plan loads its own, two plans of one shape included), which took over a second a
+// plan on a busy card, and on Windows compiles them from PTX first; and it frees and loads them all again
+// for the next plan of the same shape. A plan is never shared by two users at once - cuFFT is thread
+// safe only across distinct plans - and never destroyed, like the allocation streams above.
+//
+// A kept plan holds no work area: the user supplies one (CudaFFTPlan), so an idle plan holds only its
+// tables.
+namespace jfjoch_cufft_plans {
+    struct Plan {
+        cufftHandle handle = 0;
+        size_t work_bytes = 0;
+    };
+
+    struct Idle {
+        std::mutex m;
+        std::map<std::vector<int>, std::vector<Plan>> plans;   // by key()
+    };
+
+    inline Idle &idle() {
+        static Idle *i = new Idle;
+        return *i;
+    }
+
+    // The current device and the shape.
+    inline std::vector<int> key(const CudaFFTShape &s) {
+        int device = 0;
+        if (cudaGetDevice(&device) != cudaSuccess)
+            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read the CUDA device");
+        std::vector<int> k{device, static_cast<int>(s.n.size())};
+        k.insert(k.end(), s.n.begin(), s.n.end());
+        k.push_back(static_cast<int>(s.inembed.size()));
+        k.insert(k.end(), s.inembed.begin(), s.inembed.end());
+        k.push_back(static_cast<int>(s.onembed.size()));
+        k.insert(k.end(), s.onembed.begin(), s.onembed.end());
+        k.insert(k.end(), {s.istride, s.idist, s.ostride, s.odist, static_cast<int>(s.type), s.batch});
+        return k;
+    }
+
+    inline Plan make(const CudaFFTShape &s) {
+        Plan p;
+        cufftResult r = cufftCreate(&p.handle);
+        if (r == CUFFT_SUCCESS) {
+            r = cufftSetAutoAllocation(p.handle, 0);
+            if (r == CUFFT_SUCCESS)
+                r = cufftMakePlanMany(p.handle, static_cast<int>(s.n.size()), const_cast<int *>(s.n.data()),
+                                      s.inembed.empty() ? nullptr : const_cast<int *>(s.inembed.data()),
+                                      s.istride, s.idist,
+                                      s.onembed.empty() ? nullptr : const_cast<int *>(s.onembed.data()),
+                                      s.ostride, s.odist, s.type, s.batch, &p.work_bytes);
+            if (r != CUFFT_SUCCESS)
+                cufftDestroy(p.handle);
+        }
+        // Said as a memory failure whatever the code: on a full card (two runs sharing it) the plan
+        // failed as CUFFT_INTERNAL_ERROR, not CUFFT_ALLOC_FAILED.
+        if (r != CUFFT_SUCCESS)
+            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
+                                  GPUMemoryFailure("Creating the CUFFT plan (cufftMakePlanMany)", 0, true,
+                                                   ("cuFFT error " + std::to_string(r)).c_str()));
+        // The plan's setup can still be queued on the NULL stream when cufftMakePlanMany returns, and a
+        // non-blocking stream (CudaStream) does not wait for it: a transform launched on one straight
+        // away read the plan's tables before they were there - an illegal address in 3 of 4 runs where
+        // a rotation pool built its FFT indexers at transform length 3072. Waiting here orders it.
+        if (cudaStreamSynchronize(nullptr) != cudaSuccess)
+            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
+                                  "Creating the CUFFT plan: setup on the NULL stream failed");
+        return p;
+    }
+
+    // An idle plan of this key, or a new one.
+    inline Plan borrow(const std::vector<int> &k, const CudaFFTShape &s) {
+        {
+            std::lock_guard lock(idle().m);
+            auto it = idle().plans.find(k);
+            if (it != idle().plans.end() && !it->second.empty()) {
+                const Plan p = it->second.back();
+                it->second.pop_back();
+                return p;
+            }
+        }
+        return make(s);
+    }
+
+    inline void give_back(const std::vector<int> &k, const Plan &p) {
+        std::lock_guard lock(idle().m);
+        idle().plans[k].push_back(p);
+    }
+}
+
+// A cuFFT plan of `shape` on the current device, borrowed from the plans kept above for as long as this
+// lives, its transforms on `stream`. Its work area is `work_area` (at least WorkBytes() long), or, given
+// none, one of its own. The plan goes back once the stream has finished with it.
+class CudaFFTPlan {
+    std::vector<int> key_;
+    jfjoch_cufft_plans::Plan plan_;
+    cudaStream_t stream_ = nullptr;
+    CudaDevicePtr<char> work_;
+
+    void release() {
+        if (!plan_.handle)
+            return;
+        cudaStreamSynchronize(stream_);
+        work_ = CudaDevicePtr<char>();
+        jfjoch_cufft_plans::give_back(key_, plan_);
+        plan_ = {};
+    }
+
+public:
+    CudaFFTPlan() = default;
+
+    explicit CudaFFTPlan(const CudaFFTShape &shape, cudaStream_t stream = nullptr, void *work_area = nullptr)
+            : key_(jfjoch_cufft_plans::key(shape)), plan_(jfjoch_cufft_plans::borrow(key_, shape)), stream_(stream) {
+        try {
+            if (work_area == nullptr) {
+                work_ = CudaDevicePtr<char>(std::max<size_t>(plan_.work_bytes, 1));
+                work_area = work_;
+            }
+            if (cufftSetWorkArea(plan_.handle, work_area) != CUFFT_SUCCESS
+                || cufftSetStream(plan_.handle, stream_) != CUFFT_SUCCESS)
+                throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Setting up the CUFFT plan failed");
+        } catch (...) {
+            release();
+            throw;
+        }
+    }
+
+    // The work area a plan of `shape` on the current device needs - read from the kept plan, which is
+    // made here if there is none, rather than from a plan made only to be asked.
+    static size_t WorkBytes(const CudaFFTShape &shape) {
+        const std::vector<int> k = jfjoch_cufft_plans::key(shape);
+        const jfjoch_cufft_plans::Plan p = jfjoch_cufft_plans::borrow(k, shape);
+        jfjoch_cufft_plans::give_back(k, p);
+        return p.work_bytes;
+    }
+
+    ~CudaFFTPlan() {
+        release();
+    }
+    // Move-only type
+    CudaFFTPlan(CudaFFTPlan&& other) noexcept
+            : key_(std::move(other.key_)), plan_(other.plan_), stream_(other.stream_), work_(std::move(other.work_)) {
+        other.plan_ = {};
+    }
+    CudaFFTPlan& operator=(CudaFFTPlan&& other) noexcept {
+        if (this != &other) {
+            release();
+            key_ = std::move(other.key_);
+            plan_ = other.plan_;
+            stream_ = other.stream_;
+            work_ = std::move(other.work_);
+            other.plan_ = {};
+        }
+        return *this;
+    }
+    CudaFFTPlan(const CudaFFTPlan&) = delete;
+    CudaFFTPlan& operator=(const CudaFFTPlan&) = delete;
+
+    size_t WorkBytes() const { return plan_.work_bytes; }
+    operator cufftHandle() const { return plan_.handle; }
+    cufftHandle get() const { return plan_.handle; }
+};
+
 // Page-locked host memory is slow to get and to give back: cudaHostAlloc and cudaFreeHost map and lock
 // every page under the process's memory-map lock, so the calls of all threads queue behind each other -
 // measured at 1-3 ms a call on a workstation and 20-30 ms on a four-socket host. The engines that hold it

image_analysis/indexing/FFTIndexerGPU.cu

@@ -245,11 +245,12 @@ FFTIndexerGPU::FFTIndexerGPU(const IndexingSettings &settings)
 
     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));
+    plan = CudaFFTPlan(CudaFFTShape{.n = {static_cast<int>(histogram_size)},
+                                    .idist = static_cast<int>(histogram_size),
+                                    .odist = static_cast<int>(histogram_size / 2 + 1),
+                                    .type = CUFFT_R2C,
+                                    .batch = nDirections},
+                       stream);
 }
 
 

image_analysis/indexing/FFTIndexerGPU.h

@@ -35,10 +35,11 @@ class FFTIndexerGPU : public FFTIndexer {
 
     CudaRegisteredVector<FFTResult> result_fft_reg;
 
-    CudaFFTPlan plan;
-
     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:

image_analysis/structure_refinement/ModelStructureFactorsGPU.cu

@@ -115,34 +115,17 @@ __global__ void transpose_kernel(const float *in, int nu, int nv, int nw, float
 }
 
 // The structure factors' transform: r2c of an (nu, nv, nw) grid stored u fastest, which halves u.
-cufftHandle ForwardPlan(const ModelDensityGrid &g, size_t &work) {
-    cufftHandle plan;
-    cufft_err(cufftCreate(&plan));
-    cufft_err(cufftSetAutoAllocation(plan, 0));
-    int n[3] = {g.nw, g.nv, g.nu};
-    const int idist = g.nu * g.nv * g.nw, odist = (g.nu / 2 + 1) * g.nv * g.nw;
-    const cufftResult r = cufftMakePlanMany(plan, 3, n, nullptr, 1, idist, nullptr, 1, odist, CUFFT_R2C, 1, &work);
-    if (r != CUFFT_SUCCESS)
-        cufftDestroy(plan);
-    cufft_err(r);
-    return plan;
+CudaFFTShape ForwardShape(const ModelDensityGrid &g) {
+    return {.n = {g.nw, g.nv, g.nu}, .idist = g.nu * g.nv * g.nw, .odist = (g.nu / 2 + 1) * g.nv * g.nw,
+            .type = CUFFT_R2C};
 }
 
 // A map's transform: in-place c2r of an (nu, nv, nw) grid stored as gemmi's ZYX half-l grid, w fastest and
 // halved, the real output padded to 2 (nw / 2 + 1) per row as cuFFT does in place.
-cufftHandle MapPlan(int nu, int nv, int nw, size_t &work) {
-    cufftHandle plan;
-    cufft_err(cufftCreate(&plan));
-    cufft_err(cufftSetAutoAllocation(plan, 0));
-    int n[3] = {nu, nv, nw};
-    int inembed[3] = {nu, nv, nw / 2 + 1};
-    int onembed[3] = {nu, nv, 2 * (nw / 2 + 1)};
-    const int idist = nu * nv * (nw / 2 + 1), odist = 2 * idist;
-    const cufftResult r = cufftMakePlanMany(plan, 3, n, inembed, 1, idist, onembed, 1, odist, CUFFT_C2R, 1, &work);
-    if (r != CUFFT_SUCCESS)
-        cufftDestroy(plan);
-    cufft_err(r);
-    return plan;
+CudaFFTShape MapShape(int nu, int nv, int nw) {
+    const int idist = nu * nv * (nw / 2 + 1);
+    return {.n = {nu, nv, nw}, .inembed = {nu, nv, nw / 2 + 1}, .idist = idist,
+            .onembed = {nu, nv, 2 * (nw / 2 + 1)}, .odist = 2 * idist, .type = CUFFT_C2R};
 }
 
 // The calling thread on `device` for as long as this lives, and back on its own device after.
@@ -186,7 +169,7 @@ struct ModelStructureFactorsGPUEngineImpl {
     CudaDevicePtr<float2> f, fmask;
     std::unique_ptr<ModelDensityGPU> density;
     std::unique_ptr<ModelMaskGPU> mask;
-    cufftHandle forward = 0;
+    CudaFFTPlan forward;
 
     // On `dev`, which the caller has made current.
     ModelStructureFactorsGPUEngineImpl(int dev, const ModelStructureFactorsGPUSetup &s, size_t work)
@@ -223,20 +206,14 @@ struct ModelStructureFactorsGPUEngineImpl {
         up(row_hkl, setup.rows.data(), setup.rows.size() * 3 * sizeof(int));
         up(row_scale, setup.row_scale.data(), setup.row_scale.size() * sizeof(double));
 
-        size_t plan_work = 0;
-        forward = ForwardPlan(setup.grid, plan_work);
-        if (plan_work > work_bytes) {
-            cufftDestroy(forward);
+        forward = CudaFFTPlan(ForwardShape(setup.grid), stream, fft_work);
+        if (forward.WorkBytes() > work_bytes)
             throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "structure factors: cuFFT work area over the reserve");
-        }
-        cufft_err(cufftSetWorkArea(forward, fft_work));
-        cufft_err(cufftSetStream(forward, stream));
         cuda_err(cudaStreamSynchronize(stream));
     }
 
     ~ModelStructureFactorsGPUEngineImpl() {
         cudaStreamSynchronize(stream);
-        cufftDestroy(forward);
     }
 
     float Norm() const {
@@ -254,10 +231,7 @@ size_t ModelStructureFactorsGPUEngine::TotalMemory(int device) {
 size_t ModelStructureFactorsGPUEngine::FFTWorkBytes(int device, const ModelDensityGrid &grid, int map_nu, int map_nv,
                                                     int map_nw) {
     DeviceGuard guard(device);
-    size_t forward = 0, inverse = 0;
-    cufftDestroy(ForwardPlan(grid, forward));
-    cufftDestroy(MapPlan(map_nu, map_nv, map_nw, inverse));
-    return std::max(forward, inverse);
+    return std::max(CudaFFTPlan::WorkBytes(ForwardShape(grid)), CudaFFTPlan::WorkBytes(MapShape(map_nu, map_nv, map_nw)));
 }
 
 size_t ModelStructureFactorsGPUEngine::DeviceBytes(const ModelStructureFactorsGPUSetup &s, size_t fft_work_bytes) {
@@ -338,28 +312,19 @@ void ModelStructureFactorsGPUEngine::Map(int nu, int nv, int nw, double volume,
     const size_t complex_points = static_cast<size_t>(nu) * nv * (nw / 2 + 1);
     if (points > e.setup.map_points || complex_points > e.setup.map_complex_points)
         throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "map: grid over the reserve");
-    size_t plan_work = 0;
-    const cufftHandle plan = MapPlan(nu, nv, nw, plan_work);
-    try {
-        if (plan_work > e.work_bytes)
-            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "map: cuFFT work area over the reserve");
-        cufft_err(cufftSetWorkArea(plan, e.fft_work));
-        cufft_err(cufftSetStream(plan, e.stream));
-        cuda_err(cudaMemcpyAsync(e.spectrum, coefficients, complex_points * sizeof(float2), cudaMemcpyHostToDevice,
-                                 e.stream));
-        conjugate_kernel<<<blocks(complex_points), THREADS, 0, e.stream>>>(e.spectrum, complex_points);
-        cuda_err(cudaGetLastError());
-        cufft_err(cufftExecC2R(plan, reinterpret_cast<cufftComplex *>(e.spectrum.get()),
-                               reinterpret_cast<cufftReal *>(e.spectrum.get())));
-        transpose_kernel<<<blocks(points), THREADS, 0, e.stream>>>(reinterpret_cast<const float *>(e.spectrum.get()),
-                                                                   nu, nv, nw, static_cast<float>(1.0 / volume), e.grid);
-        cuda_err(cudaGetLastError());
-        cuda_err(cudaMemcpyAsync(map, e.grid, points * sizeof(float), cudaMemcpyDeviceToHost, e.stream));
-        cuda_err(cudaStreamSynchronize(e.stream));
-    } catch (...) {
-        cudaStreamSynchronize(e.stream);
-        cufftDestroy(plan);
-        throw;
-    }
-    cufftDestroy(plan);
+    // Its stream is waited for before the plan goes back, on a throw as well.
+    CudaFFTPlan plan(MapShape(nu, nv, nw), e.stream, e.fft_work);
+    if (plan.WorkBytes() > e.work_bytes)
+        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "map: cuFFT work area over the reserve");
+    cuda_err(cudaMemcpyAsync(e.spectrum, coefficients, complex_points * sizeof(float2), cudaMemcpyHostToDevice,
+                             e.stream));
+    conjugate_kernel<<<blocks(complex_points), THREADS, 0, e.stream>>>(e.spectrum, complex_points);
+    cuda_err(cudaGetLastError());
+    cufft_err(cufftExecC2R(plan, reinterpret_cast<cufftComplex *>(e.spectrum.get()),
+                           reinterpret_cast<cufftReal *>(e.spectrum.get())));
+    transpose_kernel<<<blocks(points), THREADS, 0, e.stream>>>(reinterpret_cast<const float *>(e.spectrum.get()),
+                                                               nu, nv, nw, static_cast<float>(1.0 / volume), e.grid);
+    cuda_err(cudaGetLastError());
+    cuda_err(cudaMemcpyAsync(map, e.grid, points * sizeof(float), cudaMemcpyDeviceToHost, e.stream));
+    cuda_err(cudaStreamSynchronize(e.stream));
 }

image_analysis/structure_refinement/MolecularReplacementGPU.cu

@@ -338,23 +338,15 @@ gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map) {
     auto d_map = upload(map.data, stream);
     CudaDevicePtr<cufftComplex> d_half(static_cast<size_t>(map.nu / 2 + 1) * map.nv * map.nw);
     CudaDevicePtr<float2> d_out(hkl.data.size());
-    cufftHandle plan;
-    cufft_err(cufftPlan3d(&plan, map.nw, map.nv, map.nu, CUFFT_R2C));
-    try {
-        cufft_err(cufftSetStream(plan, stream));
-        cufft_err(cufftExecR2C(plan, d_map.get(), d_half.get()));
-        half_u_to_half_w<<<static_cast<unsigned>((hkl.data.size() + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
-                d_half, map.nu, map.nv, map.nw, norm, d_out);
-        cuda_err(cudaGetLastError());
-        static_assert(sizeof(std::complex<float>) == sizeof(float2));
-        cuda_err(cudaMemcpyAsync(hkl.data.data(), d_out.get(), hkl.data.size() * sizeof(float2),
-                                 cudaMemcpyDeviceToHost, stream));
-        cuda_err(cudaStreamSynchronize(stream));
-    } catch (...) {
-        cufftDestroy(plan);
-        throw;
-    }
-    cufftDestroy(plan);
+    CudaFFTPlan plan(CudaFFTShape{.n = {map.nw, map.nv, map.nu}, .type = CUFFT_R2C}, stream);
+    cufft_err(cufftExecR2C(plan, d_map.get(), d_half.get()));
+    half_u_to_half_w<<<static_cast<unsigned>((hkl.data.size() + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
+            d_half, map.nu, map.nv, map.nw, norm, d_out);
+    cuda_err(cudaGetLastError());
+    static_assert(sizeof(std::complex<float>) == sizeof(float2));
+    cuda_err(cudaMemcpyAsync(hkl.data.data(), d_out.get(), hkl.data.size() * sizeof(float2),
+                             cudaMemcpyDeviceToHost, stream));
+    cuda_err(cudaStreamSynchronize(stream));
     return hkl;
 }
 
@@ -427,11 +419,7 @@ struct TranslationSetupGPU {
     CudaDevicePtr<float3> s;
     CudaDevicePtr<float2> phase;
     CudaDevicePtr<float> weight;
-    cufftHandle plan = 0;
-    ~TranslationSetupGPU() {
-        if (plan != 0)
-            cufftDestroy(plan);
-    }
+    CudaFFTPlan plan;
 };
 
 namespace {
@@ -484,7 +472,7 @@ std::unique_ptr<TranslationSetupGPU> prepare_translation(const TranslationSetup
     out->pair_k = upload(setup.pair_k, stream);
     out->offset = upload(offset, stream);
     out->entry = upload(entry, stream);
-    cufft_err(cufftPlan3d(&out->plan, nu, nv, nw, CUFFT_C2R));
+    out->plan = CudaFFTPlan(CudaFFTShape{.n = {nu, nv, nw}, .type = CUFFT_C2R}, stream);
     cuda_err(cudaStreamSynchronize(stream));   // the host vectors go out of scope
     return out;
 }

image_analysis/structure_refinement/RigidBodyGPU.cu

@@ -281,19 +281,10 @@ int blocks(size_t n) {
     return static_cast<int>((n + THREADS - 1) / THREADS);
 }
 
-// A plan for an (nu, nv, nw) R2C transform, u fastest, `batch` grids back to back, whose work area the
-// caller provides; its size in `work`.
-cufftHandle MakePlan(int nu, int nv, int nw, int batch, size_t &work) {
-    cufftHandle plan;
-    cufft_err(cufftCreate(&plan));
-    cufft_err(cufftSetAutoAllocation(plan, 0));
-    int n[3] = {nw, nv, nu};
-    const int idist = nu * nv * nw, odist = (nu / 2 + 1) * nv * nw;
-    const cufftResult r = cufftMakePlanMany(plan, 3, n, nullptr, 1, idist, nullptr, 1, odist, CUFFT_R2C, batch, &work);
-    if (r != CUFFT_SUCCESS)
-        cufftDestroy(plan);
-    cufft_err(r);
-    return plan;
+// An (nu, nv, nw) R2C transform, u fastest, `batch` grids back to back.
+CudaFFTShape GridShape(int nu, int nv, int nw, int batch) {
+    return {.n = {nw, nv, nu}, .idist = nu * nv * nw, .odist = (nu / 2 + 1) * nv * nw, .type = CUFFT_R2C,
+            .batch = batch};
 }
 
 int GreatestStreamPriority() {
@@ -351,7 +342,7 @@ struct RigidBodyGPUEngineImpl {
     CudaDevicePtr<float2> point_fcmol, point_fmask;
     CudaDevicePtr<double> products, x;
 
-    std::map<std::array<int, 4>, cufftHandle> plans;
+    std::map<std::array<int, 4>, CudaFFTPlan> plans;
     std::unique_ptr<ModelDensityGPU> density;
     std::unique_ptr<ModelMaskGPU> mask;
     std::unique_ptr<ModelScaleGPU> scale;
@@ -414,8 +405,6 @@ struct RigidBodyGPUEngineImpl {
     ~RigidBodyGPUEngineImpl() {
         cudaSetDevice(device);
         cudaStreamSynchronize(stream);
-        for (auto &[k, plan] : plans)
-            cufftDestroy(plan);
     }
 
     cufftHandle Plan(int batch) {
@@ -423,16 +412,10 @@ struct RigidBodyGPUEngineImpl {
         auto it = plans.find(k);
         if (it != plans.end())
             return it->second;
-        size_t work = 0;
-        const cufftHandle plan = MakePlan(geom.nu, geom.nv, geom.nw, batch, work);
-        if (work > cap.fft_work_bytes) {
-            cufftDestroy(plan);
+        CudaFFTPlan plan(GridShape(geom.nu, geom.nv, geom.nw, batch), stream, fft_work);
+        if (plan.WorkBytes() > cap.fft_work_bytes)
             throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: cuFFT work area over the reserve");
-        }
-        cufft_err(cufftSetWorkArea(plan, fft_work));
-        cufft_err(cufftSetStream(plan, stream));
-        plans.emplace(k, plan);
-        return plan;
+        return plans.emplace(k, std::move(plan)).first->second;
     }
 
     size_t GridPoints() const { return static_cast<size_t>(geom.nu) * geom.nv * geom.nw; }
@@ -496,12 +479,8 @@ int RigidBodyGPUEngine::CurrentDevice() {
 
 size_t RigidBodyGPUEngine::FFTWorkBytes(int nu, int nv, int nw) {
     size_t largest = 0;
-    for (int batch : {1, 3}) {
-        size_t work = 0;
-        const cufftHandle plan = MakePlan(nu, nv, nw, batch, work);
-        cufftDestroy(plan);
-        largest = std::max(largest, work);
-    }
+    for (int batch : {1, 3})
+        largest = std::max(largest, CudaFFTPlan::WorkBytes(GridShape(nu, nv, nw, batch)));
     return largest;
 }
 
@@ -606,6 +585,8 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) {
     e.mask->SetGrid(mask_grid, images);
     e.density->SetAtoms(DensityGrid(zone), zone.atoms);
     e.scale->SetPoints(zone.point_hkl, zone.point_stol2, zone.point_fobs, zone.point_sigma, zone.constraints, zone.frac);
+    // The previous zone's plans go back to be lent to another engine; this one takes the zone's own.
+    e.plans.clear();
     e.Plan(1);
     e.Plan(3);
     cuda_err(cudaStreamSynchronize(e.stream));

tests/CUDAMemHelpersTest.cpp

@@ -9,6 +9,7 @@
 
 #ifdef JFJOCH_USE_CUDA
 
+#include <cstring>
 #include <future>
 #include <latch>
 #include <set>
@@ -128,6 +129,52 @@ TEST_CASE("CudaPinnedPool_ReusedByTheNextEngine", "[CUDAMemHelpers]") {
     CHECK(cuda_release_pinned_pool() == 0);
 }
 
+// A cuFFT plan is kept once made and lent to one user at a time: two users at once get two plans, the
+// next one after them gets one of those back, and a transform on a lent plan - another stream, a work
+// area of the caller's - gives the same result bit for bit.
+TEST_CASE("CudaFFTPlan_KeptAndLentToOneUserAtATime", "[CUDAMemHelpers]") {
+    if (get_gpu_count() == 0)
+        SKIP("No CUDA GPU present");
+
+    const int n = 1000, batch = 7, nc = n / 2 + 1;
+    const CudaFFTShape shape{.n = {n}, .idist = n, .odist = nc, .type = CUFFT_R2C, .batch = batch};
+    std::vector<float> input(n * batch);
+    for (size_t i = 0; i < input.size(); i++)
+        input[i] = static_cast<float>((i * 7919) % 1013) - 500.0f;
+    CudaDevicePtr<float> d_in(input.size());
+    CudaDevicePtr<cufftComplex> d_out(nc * batch);
+    REQUIRE(cudaMemcpy(d_in, input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice) == cudaSuccess);
+
+    auto transform = [&](const CudaFFTPlan &plan, cudaStream_t stream) {
+        REQUIRE(cufftExecR2C(plan, d_in, d_out) == CUFFT_SUCCESS);
+        std::vector<cufftComplex> out(nc * batch);
+        REQUIRE(cudaMemcpyAsync(out.data(), d_out, out.size() * sizeof(cufftComplex), cudaMemcpyDeviceToHost,
+                                stream) == cudaSuccess);
+        REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
+        return out;
+    };
+
+    CudaStream stream_a, stream_c;
+    std::vector<cufftComplex> first;
+    cufftHandle handle_a = 0;
+    {
+        CudaFFTPlan a(shape, stream_a);
+        CudaFFTPlan b(shape);
+        handle_a = a.get();
+        CHECK(handle_a != b.get());
+        first = transform(a, stream_a);
+    }
+
+    const size_t work = CudaFFTPlan::WorkBytes(shape);
+    CudaDevicePtr<char> work_area(std::max<size_t>(work, 1));
+    CudaFFTPlan c(shape, stream_c, work_area);
+    CHECK(c.get() == handle_a);
+    CHECK(c.WorkBytes() == work);
+    const std::vector<cufftComplex> again = transform(c, stream_c);
+    REQUIRE(again.size() == first.size());
+    CHECK(memcmp(again.data(), first.data(), first.size() * sizeof(cufftComplex)) == 0);
+}
+
 // An FFT indexer that finds the card full waits for the GPU work beside it to give its memory back, and
 // is built then: the main line's for the probes, a probe's for the main line's image pass - never for a
 // probe, itself included. With nothing to wait for it fails at once, as one that does not fit always has.