PERF: Cache GPU context, plan, and buffers across transforms; add leak test and export VKFFT_BACKEND

CMakeLists.txt

@@ -131,7 +131,6 @@ elseif(VKFFT_BACKEND EQUAL 5)
     GIT_SHALLOW TRUE
   )
   FetchContent_MakeAvailable(metal_cpp)
-  list(APPEND VkFFTBackend_SYSTEM_INCLUDE_DIRS ${metal_cpp_SOURCE_DIR})
 
   find_library(METAL_FRAMEWORK Metal REQUIRED)
   find_library(FOUNDATION_FRAMEWORK Foundation REQUIRED)
@@ -236,7 +235,6 @@ else()
   )
   FetchContent_MakeAvailable(vkfft_header_only)
 endif()
-list(APPEND VkFFTBackend_SYSTEM_INCLUDE_DIRS ${vkfft_INCLUDE_DIR})
 
 #### Configure ITK module ####
 
@@ -249,6 +247,56 @@ else()
   itk_module_impl()
 endif()
 
+# VkCommon's ABI layout is selected by VKFFT_BACKEND, so consumers must compile
+# itkVkCommon.h with the same value the library was built with; export it.
+target_compile_definitions(VkFFTBackend PUBLIC VKFFT_BACKEND=${VKFFT_BACKEND})
+
+# itkVkCommon.h includes the fetched vkFFT/metal-cpp headers. Carry their
+# build-tree locations on the target as a build-only property: BUILD_INTERFACE
+# keeps them out of install(EXPORT) so the exported target stays relocatable,
+# while installed consumers resolve the headers from ITK_INSTALL_INCLUDE_DIR
+# where they are installed below.
+target_include_directories(
+  VkFFTBackend
+  SYSTEM
+  PUBLIC
+    "$<BUILD_INTERFACE:${vkfft_INCLUDE_DIR}>"
+)
+if(VKFFT_BACKEND EQUAL 5)
+  target_include_directories(
+    VkFFTBackend
+    SYSTEM
+    PUBLIC
+      "$<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>"
+  )
+endif()
+
+# itkVkCommon.h includes these fetched headers, so install them for consumers.
+install(
+  DIRECTORY
+    "${vkfft_INCLUDE_DIR}/"
+  DESTINATION "${ITK_INSTALL_INCLUDE_DIR}"
+  COMPONENT Development
+  FILES_MATCHING
+  PATTERN
+  "*.h"
+  PATTERN
+  "*.hpp"
+)
+if(VKFFT_BACKEND EQUAL 5)
+  install(
+    DIRECTORY
+      "${metal_cpp_SOURCE_DIR}/Foundation"
+      "${metal_cpp_SOURCE_DIR}/Metal"
+      "${metal_cpp_SOURCE_DIR}/QuartzCore"
+    DESTINATION "${ITK_INSTALL_INCLUDE_DIR}"
+    COMPONENT Development
+    FILES_MATCHING
+    PATTERN
+    "*.hpp"
+  )
+endif()
+
 if(VKFFT_BACKEND EQUAL 4)
   target_link_libraries(VkFFTBackend PUBLIC ${LevelZero_LIBRARY})
   target_include_directories(

include/itkVkCommon.h

@@ -109,6 +109,19 @@ class VkFFTBackend_EXPORT VkCommon
              this->inputBufferBytes != rhs.inputBufferBytes || this->outputCPUBuffer != rhs.outputCPUBuffer ||
              this->outputBufferBytes != rhs.outputBufferBytes;
     }
+
+    /** Compare only the transform-shape fields, ignoring the per-call CPU buffer
+     * pointers and byte counts, which change on every call. The VkFFT plan depends
+     * only on the shape, so a cached plan is reusable across calls that differ
+     * only in their buffers. */
+    bool
+    SameShapeAs(const VkParameters & rhs) const
+    {
+      return this->X == rhs.X && this->Y == rhs.Y && this->Z == rhs.Z && this->P == rhs.P && this->B == rhs.B &&
+             this->N == rhs.N && this->fft == rhs.fft && this->PSize == rhs.PSize && this->I == rhs.I &&
+             this->normalized == rhs.normalized && this->omitDimension[0] == rhs.omitDimension[0] &&
+             this->omitDimension[1] == rhs.omitDimension[1] && this->omitDimension[2] == rhs.omitDimension[2];
+    }
   };
 
   struct VkGPU
@@ -181,10 +194,38 @@ class VkFFTBackend_EXPORT VkCommon
   VkParameters       m_VkParameters{};
   VkFFTConfiguration m_VkFFTConfiguration{};
 
-  // Re-create GPU kernel if these members indicate to
-  bool         m_MustConfigure{ true };
-  VkGPU        m_VkGPUPrevious{};
-  VkParameters m_VkParametersPrevious{};
+  // Cached GPU context + plan configuration are (re)built only on first use or when
+  // the device/transform shape changes; m_VkGPU and m_VkParameters hold that cached state.
+  bool m_MustConfigure{ true };
+
+  // Cached compiled plan and persistent per-shape GPU buffers. initializeVkFFT (which
+  // JIT-compiles the FFT kernels) and the buffer allocations run once per shape and are
+  // reused across same-shape transforms; per call only the host<->device copies and the
+  // VkFFTAppend run. Released together with the context in ReleaseBackend().
+  //
+  // The VkFFTApplication is heap-allocated (not an inline member) so that its size is
+  // computed in the library translation unit that actually populates it; embedding it
+  // by value makes the class layout depend on sizeof(VkFFTApplication) at every include
+  // site, which can differ and corrupt the members that follow.
+  VkFFTApplication * m_VkFFTApplication{ nullptr };
+  bool               m_PlanConfigured{ false };
+#if (VKFFT_BACKEND == CUDA)
+  cuFloatComplex * m_GPUBuffer{ nullptr };
+  cuFloatComplex * m_InputGPUBuffer{ nullptr };
+  cuFloatComplex * m_OutputGPUBuffer{ nullptr };
+#elif (VKFFT_BACKEND == OPENCL)
+  cl_mem m_GPUBuffer{ nullptr };
+  cl_mem m_InputGPUBuffer{ nullptr };
+  cl_mem m_OutputGPUBuffer{ nullptr };
+#elif (VKFFT_BACKEND == LEVEL_ZERO)
+  void * m_GPUBuffer{ nullptr };
+  void * m_InputGPUBuffer{ nullptr };
+  void * m_OutputGPUBuffer{ nullptr };
+#elif (VKFFT_BACKEND == METAL)
+  MTL::Buffer * m_GPUBuffer{ nullptr };
+  MTL::Buffer * m_InputGPUBuffer{ nullptr };
+  MTL::Buffer * m_OutputGPUBuffer{ nullptr };
+#endif
 };
 
 } // namespace itk