diff --git a/image_analysis/geom_refinement/BeamCenterFFTGPU.cu b/image_analysis/geom_refinement/BeamCenterFFTGPU.cu index 98a7e8a88..d54e2fcf6 100644 --- a/image_analysis/geom_refinement/BeamCenterFFTGPU.cu +++ b/image_analysis/geom_refinement/BeamCenterFFTGPU.cu @@ -45,6 +45,33 @@ void Check(cufftResult res, const char *what) { std::vector Dims(int64_t a) { return {static_cast(a)}; } std::vector Dims(int64_t a, int64_t b) { return {static_cast(a), static_cast(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(ny * nx), + .onembed = Dims(ny, nxc), .odist = static_cast(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(ny * nxc), + .onembed = Dims(ny, nx), .odist = static_cast(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(nfft), + .onembed = Dims(nc), .odist = static_cast(nc), .type = CUFFT_R2C, + .batch = static_cast(nbatch)}; +} + +CudaFFTShape LineInverse(int64_t nfft) { + const int64_t nc = nfft / 2 + 1; + return {.n = Dims(nfft), .inembed = Dims(nc), .idist = static_cast(nc), + .onembed = Dims(nfft), .odist = static_cast(nfft), .type = CUFFT_C2R, .batch = 4}; +} + __global__ void SquareInPlace(float *data, int64_t n) { const int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; if (i < n) @@ -249,10 +276,8 @@ PointSurfacesDevice TransformPoint(const std::vector &a, const std::vecto out.real = CudaDevicePtr(nreal, ALLOC); CudaDevicePtr d_a(ncomplex, ALLOC), d_m(ncomplex, ALLOC), d_prod(ncomplex, ALLOC); - CudaFFTPlan fwd(2, Dims(ny, nx), Dims(ny, nx), 1, static_cast(nreal), Dims(ny, nxc), 1, - static_cast(ncomplex), CUFFT_R2C, 1); - CudaFFTPlan inv(2, Dims(ny, nx), Dims(ny, nxc), 1, static_cast(ncomplex), Dims(ny, nx), 1, - static_cast(nreal), CUFFT_C2R, 1); + CudaFFTPlan fwd(PointForward(ny, nx)); + CudaFFTPlan inv(PointInverse(ny, nx)); auto upload = [&](const std::vector &src) { Check(cudaMemset(out.real.get(), 0, nreal * sizeof(float)), @@ -428,8 +453,7 @@ BeamCenterConvSurfaces1D BeamCenterFFTGPU::LineSurfaces(const std::vector { CudaDevicePtr d_img(h * w, ALLOC), d_in(nfft * nbatch, ALLOC); CudaDevicePtr 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(nfft), Dims(nc), 1, - static_cast(nc), CUFFT_R2C, static_cast(nbatch)); + CudaFFTPlan fwd(LineForward(nfft, nbatch)); auto transform = [&](int square, cufftComplex *dst) { GatherSequences<<>>(d_img.get(), nseq, nbatch, stride_seq, @@ -459,8 +483,7 @@ BeamCenterConvSurfaces1D BeamCenterFFTGPU::LineSurfaces(const std::vector CudaDevicePtr d_out(4 * nfft, ALLOC); PackAccumulators<<>>(accC, accS, accQ, accD, nc, d_pack.get()); CheckLastKernel("BeamCenterFFT: accumulator packing failed"); - CudaFFTPlan inv(1, Dims(nfft), Dims(nc), 1, static_cast(nc), Dims(nfft), 1, - static_cast(nfft), CUFFT_C2R, 4); + CudaFFTPlan inv(LineInverse(nfft)); Check(cufftExecC2R(inv, d_pack, d_out), "BeamCenterFFT: inverse transform failed"); std::vector host(static_cast(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(ny), static_cast(nx), CUFFT_R2C, &work_fwd); - cufftEstimate2d(static_cast(ny), static_cast(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 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(nfft), - embed_complex.data(), 1, static_cast(nc), CUFFT_R2C, - static_cast(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; } diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index 6d2db7ad0..6de2287d9 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/image_analysis/indexing/CUDAMemHelpers.h @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -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(n), - const_cast(inembed), istride, idist, - const_cast(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& n, - const std::vector& inembed, int istride, int idist, - const std::vector& 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 n; + std::vector inembed; + int istride = 1, idist = 0; + std::vector 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> plans; // by key() + }; + + inline Idle &idle() { + static Idle *i = new Idle; + return *i; + } + + // The current device and the shape. + inline std::vector key(const CudaFFTShape &s) { + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read the CUDA device"); + std::vector k{device, static_cast(s.n.size())}; + k.insert(k.end(), s.n.begin(), s.n.end()); + k.push_back(static_cast(s.inembed.size())); + k.insert(k.end(), s.inembed.begin(), s.inembed.end()); + k.push_back(static_cast(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(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(s.n.size()), const_cast(s.n.data()), + s.inembed.empty() ? nullptr : const_cast(s.inembed.data()), + s.istride, s.idist, + s.onembed.empty() ? nullptr : const_cast(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 &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 &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 key_; + jfjoch_cufft_plans::Plan plan_; + cudaStream_t stream_ = nullptr; + CudaDevicePtr work_; + + void release() { + if (!plan_.handle) + return; + cudaStreamSynchronize(stream_); + work_ = CudaDevicePtr(); + 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(std::max(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 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 diff --git a/image_analysis/indexing/FFTIndexerGPU.cu b/image_analysis/indexing/FFTIndexerGPU.cu index 418b06f00..d86560d9f 100644 --- a/image_analysis/indexing/FFTIndexerGPU.cu +++ b/image_analysis/indexing/FFTIndexerGPU.cu @@ -245,11 +245,12 @@ FFTIndexerGPU::FFTIndexerGPU(const IndexingSettings &settings) DirectionsChanged(); - int n[1] = {static_cast(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(histogram_size)}, + .idist = static_cast(histogram_size), + .odist = static_cast(histogram_size / 2 + 1), + .type = CUFFT_R2C, + .batch = nDirections}, + stream); } diff --git a/image_analysis/indexing/FFTIndexerGPU.h b/image_analysis/indexing/FFTIndexerGPU.h index 4f4938545..7b5fe0e2f 100644 --- a/image_analysis/indexing/FFTIndexerGPU.h +++ b/image_analysis/indexing/FFTIndexerGPU.h @@ -35,10 +35,11 @@ class FFTIndexerGPU : public FFTIndexer { CudaRegisteredVector 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, size_t nspots) override; void DirectionsChanged() override; public: diff --git a/image_analysis/structure_refinement/ModelStructureFactorsGPU.cu b/image_analysis/structure_refinement/ModelStructureFactorsGPU.cu index ce4b7eb94..aa1b036d0 100644 --- a/image_analysis/structure_refinement/ModelStructureFactorsGPU.cu +++ b/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 f, fmask; std::unique_ptr density; std::unique_ptr 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(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<<>>(e.spectrum, complex_points); - cuda_err(cudaGetLastError()); - cufft_err(cufftExecC2R(plan, reinterpret_cast(e.spectrum.get()), - reinterpret_cast(e.spectrum.get()))); - transpose_kernel<<>>(reinterpret_cast(e.spectrum.get()), - nu, nv, nw, static_cast(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<<>>(e.spectrum, complex_points); + cuda_err(cudaGetLastError()); + cufft_err(cufftExecC2R(plan, reinterpret_cast(e.spectrum.get()), + reinterpret_cast(e.spectrum.get()))); + transpose_kernel<<>>(reinterpret_cast(e.spectrum.get()), + nu, nv, nw, static_cast(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)); } diff --git a/image_analysis/structure_refinement/MolecularReplacementGPU.cu b/image_analysis/structure_refinement/MolecularReplacementGPU.cu index d611966e8..2860a6617 100644 --- a/image_analysis/structure_refinement/MolecularReplacementGPU.cu +++ b/image_analysis/structure_refinement/MolecularReplacementGPU.cu @@ -338,23 +338,15 @@ gemmi::FPhiGrid MapToFPhiGPU(const gemmi::Grid &map) { auto d_map = upload(map.data, stream); CudaDevicePtr d_half(static_cast(map.nu / 2 + 1) * map.nv * map.nw); CudaDevicePtr 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<<((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) == 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<<((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) == 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 s; CudaDevicePtr phase; CudaDevicePtr weight; - cufftHandle plan = 0; - ~TranslationSetupGPU() { - if (plan != 0) - cufftDestroy(plan); - } + CudaFFTPlan plan; }; namespace { @@ -484,7 +472,7 @@ std::unique_ptr 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; } diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cu b/image_analysis/structure_refinement/RigidBodyGPU.cu index b5a40a784..5005d3f18 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cu +++ b/image_analysis/structure_refinement/RigidBodyGPU.cu @@ -281,19 +281,10 @@ int blocks(size_t n) { return static_cast((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 point_fcmol, point_fmask; CudaDevicePtr products, x; - std::map, cufftHandle> plans; + std::map, CudaFFTPlan> plans; std::unique_ptr density; std::unique_ptr mask; std::unique_ptr 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(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)); diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index c349d9d88..8c3769ac4 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -9,6 +9,7 @@ #ifdef JFJOCH_USE_CUDA +#include #include #include #include @@ -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 input(n * batch); + for (size_t i = 0; i < input.size(); i++) + input[i] = static_cast((i * 7919) % 1013) - 500.0f; + CudaDevicePtr d_in(input.size()); + CudaDevicePtr 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 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 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 work_area(std::max(work, 1)); + CudaFFTPlan c(shape, stream_c, work_area); + CHECK(c.get() == handle_a); + CHECK(c.WorkBytes() == work); + const std::vector 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.