PERF: Cache GPU context, plan, and buffers across transforms; add leak test and export VKFFT_BACKEND - #83
Conversation
VkCommon::Run() overwrote the cached VkGPU (clearing the live context handle) before ReleaseBackend() could free it, and the reconfigure guard compared VkParameters including the per-call CPU buffer pointers, so every Update() tore down and recreated the device context. The cleared handle meant ReleaseBackend() freed nothing, so each transform leaked one GPU context/command queue and device memory grew until clCreateContext (or cuCtxCreate) failed with an out-of-memory error after a few dozen calls. Reconfigure only when the requested device or the transform shape changes (a new SameShapeAs() comparison that ignores the buffer pointers), and keep the cached context for the lifetime of the VkCommon, released once in the destructor. ReleaseBackend() now nulls the freed CUDA/OpenCL handles so a later shape-change reconfigure cannot double-free (Level Zero and Metal already did this). The change is in the backend-agnostic Run(), so it fixes CUDA, OpenCL, Level Zero, and Metal in one place. Besides removing the leak it eliminates the per-call context creation: VkFFT becomes transform-size-dependent (compute-bound) instead of a fixed ~220 ms/call. Verified on an RTX 6000 Ada: OpenCL 256^3 222 -> 37 ms, CUDA 256^3 316 -> 147 ms, round-trip error 0.000, device memory flat across 200 iterations. Add itkVkFFTRepeatedTransformLeakTest, which reuses one filter for 200 forward+inverse transforms. It fails at ~iteration 54 with the original out-of-memory abort and passes to completion with the fix.
Building on the context caching, this also reuses the compiled VkFFT plan and the device buffers across same-shape transforms. Previously every PerformFFT allocated GPU buffers and called initializeVkFFT -- which JIT-compiles the FFT kernels (nvrtc on CUDA, clBuildProgram on OpenCL) -- then freed them again, so the per-call cost was dominated by (re)compilation and allocation rather than the transform. The VkFFTApplication and the three device buffers are now members, built once per shape (guarded by m_PlanConfigured) and released in ReleaseBackend() on the next device/shape change or at destruction. Each call only copies the host buffers and runs VkFFTAppend, with the buffers bound per call via VkFFTLaunchParams. Distinct, non-aliased buffers are freed exactly once. Two correctness points the caching exposed: - CUDA runtime calls act on the current context; with several filters (each its own VkCommon and context) the current one belongs to whichever configured last, so PerformFFT and ReleaseBackend now make this instance's context current before touching its cached buffers/plan. - VkFFTApplication is heap-allocated rather than an inline member, so its size is computed in the library TU that populates it; embedding it by value makes the class layout depend on sizeof(VkFFTApplication) at every include site, which can differ and corrupt the following members. Verified on an RTX 6000 Ada (OpenCL and CUDA), compute-sanitizer clean, leak regression test passing: VkFFT is now compute-bound and the CUDA/OpenCL results converge (e.g. 256^3 ~20 ms both, vs 222 ms OpenCL / 316 ms CUDA originally). Level Zero and Metal carry the same structural change but were not exercised on this hardware.
… ABI VkCommon embeds backend-selected members (VkGPU handles, per-shape GPU buffers) whose type and size depend on VKFFT_BACKEND, and the templated Vk*FFTImageFilter holds a VkCommon by value. The backend was set only via add_compile_definitions(), which is build-tree scoped and never reaches the target's INTERFACE_COMPILE_DEFINITIONS. An installed-module consumer thus compiled itkVkCommon.h without VKFFT_BACKEND, fell back to the OpenCL default in itkVkDefinitions.h, and built a VkCommon whose member offsets disagreed with the library. With plan/buffer caching the library dereferences a cached member pointer in ReleaseBackend(), so the mismatch read a null m_VkFFTApplication and crashed (EXC_BAD_ACCESS) on the first transform. Export the selection on the target so every consumer compiles the header with the same backend the library was built with. Verified: a default external build (no -DVKFFT_BACKEND) now runs the Metal benchmark to completion, round-trip error 0.000.
hjmjohnson
changed the title
|
Re-validated on Apple Silicon (macOS) built from this PR's current tip ( Latest Metal measurements (VkFFT vs CPU backends, 3D float)Mean ms per forward+inverse, default external build:
256³ is the 30-iter, 3-rep steady-state (PocketFFT 12.4/12.5/12.8, VkFFT 8.48/8.83/8.53); the other rows are 10-iter. Round-trip error 0.000 at every size. VkFFT-Metal beats the PocketFFT CPU backend for sizes ≥192³. Leak regression test on Metal
|
|
Added Two coupled install defects this fixesA — B — fetched headers were never installed (the deeper, latent one). Verification (macOS / Metal, no ITK core patch)
|
|
Just an FYI, there are two approaches to these type of third party libraries with regard to the cmake interface. 1) The library, headers and interface can be exported as part of the monolithic ITK exports, this approach should use the itk modular macros and involves duplicating the install and export code of the library. 2) The third-party library itself can export as a separate package/project and be imported in the exports code. Neither approach is without goblins and orcs so be careful and what ever is easiest and works is good. P.S. Using file sets can make the install and include paths easier as they also become a property of the target and not separate cmake code. |
|
@dzenanz @blowekamp FYI: I did a lot of testing for ANTs' use of this on "256^3" images, and for those usage patterns, transfer + compute time, the default PocketFFT was always the fastest; the GPU variants could get close, but never beat the "transfer + compute" times. PocketFFT ~0 transfer time; GPU very small compute time. |
itkVkCommon.h is an installed PUBLIC header that includes the fetched vkFFT.h and (on Apple) the metal-cpp headers, but those headers lived only in the build tree's _deps/, so two install defects broke consuming this module from a relocated ITK install (e.g. ANTs against installed ITKv6): 1. The fetched build-tree include dirs were added to VkFFTBackend_SYSTEM_INCLUDE_DIRS, which ITKModuleMacros records verbatim in install(EXPORT ITKTargets); CMake >= 4.x then fails generation with "INTERFACE_INCLUDE_DIRECTORIES ... prefixed in the build directory." 2. The fetched headers were never installed, so a relocated install failed with "vkFFT.h / Foundation.hpp file not found." Fix: install the fetched vkFFT (all backends) and metal-cpp (Metal) headers into ITK_INSTALL_INCLUDE_DIR alongside itkVkCommon.h, where every ITK consumer already searches. Carry the build-tree locations on the VkFFTBackend target as a build-only property (target_include_directories SYSTEM PUBLIC $<BUILD_INTERFACE:...>): the in-tree library, test driver, and header test resolve the headers at build time, while BUILD_INTERFACE keeps the build path out of install(EXPORT) so the exported target stays relocatable. Using a target property rather than VkFFTBackend_SYSTEM_GENEX_INCLUDE_DIRS also works on ITK 5.4, which has no such variable, per blowekamp's suggestion to make includes a property of the target. Verified building and running (FFT round-trip + repeated-transform leak test) against both ITK 5.4 and ITK 6.0 with VKFFT_BACKEND=3 (OpenCL), and a DESTDIR install carries no build-tree paths in the exported config.
fffc8a6 to
1e610ac
Compare
|
Thanks @blowekamp — took your steer to make the fetched-header include a property of the target rather than separate include-dir variables. The relocatable rewrite of Local verification (ITK 5.4.0 + ITK 6.0, OpenCL)Reproduced the
|
4 participants
Makes VkFFT usable for iterative workloads (e.g. registration in ANTsX/ANTs#1331) by caching the GPU context, the compiled VkFFT plan, and the per-shape device buffers across transforms, and by exporting the backend selection so out-of-tree consumers link an ABI-compatible library. Three commits, each independently reviewable:
PERFcontext caching — fixes a per-Update()context leak and removes per-call context creation.PERFplan + buffer caching — keeps the compiled plan and GPU buffers alive across same-shape calls; per call only the host↔device copy +VkFFTAppendrun.COMPexportVKFFT_BACKEND— fixes a SIGSEGV that hit every out-of-tree consumer.1 — Context-leak root cause
VkCommon::Run()didm_VkGPU = vkGPU;at the top — overwriting the cachedVkGPU(and its live context handle) with the caller's empty one — before callingReleaseBackend(). SoReleaseBackend()saw a null handle and freed nothing, while the previous call's context was abandoned. The reconfigure guard also comparedVkParametersviaoperator!=, which includes the per-call CPU buffer pointers, so it fired on every call. Net effect: every transform created a new device context/command queue and leaked the previous one; device memory grew untilclCreateContextreturned-6(CL_OUT_OF_HOST_MEMORY) / VkFFT4045(orcuCtxCreatefailed) after ~50 calls.Fix:
Run()(outside any#if VKFFT_BACKEND) now reconfigures only when the device or the transform shape changes, via a newVkParameters::SameShapeAs()that ignores the buffer pointers/sizes.ReleaseBackend()nulls the freed CUDA/OpenCL handles so a genuine shape-change reconfigure cannot double-free. Because the fix is in the shared dispatcher, it covers CUDA, OpenCL, Level Zero, and Metal at once.2 — Plan + buffer caching
The compiled
VkFFTApplicationand the per-shape GPU buffers are now cached members, built once per shape and released inReleaseBackend(). Each call only copies host↔device and runsVkFFTAppend(buffers bound per-call viaVkFFTLaunchParams). This removes the per-call plan build — on CUDA that was the nvrtc JIT; on OpenCL the (cached)clBuildProgram— so CUDA and OpenCL converge to compute-bound timings.Two non-obvious traps fixed (detail in the commit message): (a) CUDA runtime calls act on the current context, so with separate forward/inverse filters each owning a context a cached buffer became invalid under the other's context — fixed by
cuCtxSetCurrentper instance; (b) embeddingVkFFTApplicationby value madeVkCommon's layout depend onsizeof(VkFFTApplication)at every include site — fixed by heap-allocating it.3 — Export VKFFT_BACKEND (fixes out-of-tree SIGSEGV)
VkCommonembeds backend-selected members (VkGPUhandles, per-shape GPU buffers) whose type and size depend onVKFFT_BACKEND, and the templatedVk*FFTImageFilterholds aVkCommonby value. The backend was set only viaadd_compile_definitions(), which is build-tree scoped and never reaches the target'sINTERFACE_COMPILE_DEFINITIONS. An installed-module consumer therefore compileditkVkCommon.hwithoutVKFFT_BACKEND, hit the defensive#define VKFFT_BACKEND OPENCLfallback initkVkDefinitions.h, and built aVkCommonwhose member offsets disagreed with the library. With plan/buffer caching the library dereferences a cached member pointer inReleaseBackend(), so the mismatch read a nullm_VkFFTApplicationand crashed (EXC_BAD_ACCESS) on the first transform.Fix:
target_compile_definitions(VkFFTBackend PUBLIC VKFFT_BACKEND=${VKFFT_BACKEND})so every consumer compiles the header with the same backend the library was built with. This is the path ANTs (#1331) uses; the in-tree ctest never exercised it because it compiles with the module's private define.Results — VkFFT 256³ (mean ms)
Round-trip error 0.000 throughout, including across shape changes. With both caching layers CUDA and OpenCL converge (compute-bound, plan reuse removes the JIT asymmetry). On Apple Silicon unified memory the original per-call cost is far below the discrete-GPU numbers (no host↔device copy); the full caching stack drops VkFFT 256³ from 29 ms to 8.7 ms — fast enough that VkFFT-Metal beats the PocketFFT CPU backend for sizes ≥192³. (Metal context-only was not separately measured; the
—cell reflects that, not a regression.)Regression test
itkVkFFTRepeatedTransformLeakTestreuses a single forward+inverse filter pair for 200 transforms (the pattern an iterative registration uses) and checks a constant round-trips. Fails at ~iteration 54 (out-of-memory abort) on the pre-fix code and passes all 200 on the fixed code (ctest, 2.65 s on OpenCL). Also built and run on Metal (Apple Silicon): passes all 200 with flat device memory. Built and run locally before submitting.