diff --git a/CMakeLists.txt b/CMakeLists.txt index c7e7b56..9c704c2 100644 --- a/CMakeLists.txt +++ b/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 + "$" +) +if(VKFFT_BACKEND EQUAL 5) + target_include_directories( + VkFFTBackend + SYSTEM + PUBLIC + "$" + ) +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( diff --git a/include/itkVkCommon.h b/include/itkVkCommon.h index a09fe6b..5762ad4 100644 --- a/include/itkVkCommon.h +++ b/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 diff --git a/src/itkVkCommon.cxx b/src/itkVkCommon.cxx index 5030620..e70de41 100644 --- a/src/itkVkCommon.cxx +++ b/src/itkVkCommon.cxx @@ -40,15 +40,21 @@ VkCommon::Run(const VkGPU & vkGPU, const VkParameters & vkParameters) { VkFFTResult resFFT{ VKFFT_SUCCESS }; - m_VkGPU = vkGPU; - m_VkParameters = vkParameters; - if (m_MustConfigure || m_VkGPU != m_VkGPUPrevious || m_VkParameters != m_VkParametersPrevious) + // Create the GPU context once and reuse it across calls. Reconfigure only when the + // requested device or the transform shape changes -- never merely because the CPU + // buffer pointers differ (they change on every call). The cached context is released + // in the destructor, so transforms no longer create (and leak) a context per call. + const bool deviceChanged{ vkGPU.device_id != m_VkGPU.device_id }; + const bool shapeChanged{ !vkParameters.SameShapeAs(m_VkParameters) }; + if (m_MustConfigure || deviceChanged || shapeChanged) { resFFT = this->ReleaseBackend(); if (resFFT != VKFFT_SUCCESS) { return resFFT; } + m_VkGPU.device_id = vkGPU.device_id; + m_VkParameters = vkParameters; resFFT = this->ConfigureBackend(); if (resFFT != VKFFT_SUCCESS) { @@ -56,6 +62,15 @@ VkCommon::Run(const VkGPU & vkGPU, const VkParameters & vkParameters) } this->m_MustConfigure = false; } + else + { + // Cached context and plan configuration are still valid; only the per-call CPU + // buffer pointers and byte counts differ. + m_VkParameters.inputCPUBuffer = vkParameters.inputCPUBuffer; + m_VkParameters.inputBufferBytes = vkParameters.inputBufferBytes; + m_VkParameters.outputCPUBuffer = vkParameters.outputCPUBuffer; + m_VkParameters.outputBufferBytes = vkParameters.outputBufferBytes; + } resFFT = this->PerformFFT(); if (resFFT != VKFFT_SUCCESS) @@ -417,64 +432,62 @@ VkCommon::PerformFFT() #if (VKFFT_BACKEND == CUDA) cudaError resCu{ cudaSuccess }; - cuFloatComplex * inputGPUBuffer{ nullptr }; - cuFloatComplex * GPUBuffer{ nullptr }; - cuFloatComplex * outputGPUBuffer{ nullptr }; + // Make this instance's context current. With several filters (each its own VkCommon and + // CUDA context), the current context otherwise belongs to whichever filter configured last, + // and the runtime API would touch this instance's cached buffers under the wrong context. + cuCtxSetCurrent(m_VkGPU.context); - // Allocate the in-place-computation buffer - const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; - resCu = cudaMalloc((void **)&GPUBuffer, bufferBytes); - - if (resCu != cudaSuccess) + if (!m_PlanConfigured) { - std::cerr << __FILE__ "(" << __LINE__ << "): cudaMalloc returned " << resCu << std::endl; - return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - } - - m_VkFFTConfiguration.buffer = reinterpret_cast(&GPUBuffer); + // Allocate the in-place-computation buffer (persistent for this shape). + const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; + resCu = cudaMalloc((void **)&m_GPUBuffer, bufferBytes); + if (resCu != cudaSuccess) + { + std::cerr << __FILE__ "(" << __LINE__ << "): cudaMalloc returned " << resCu << std::endl; + return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; + } + m_VkFFTConfiguration.buffer = reinterpret_cast(&m_GPUBuffer); - if (m_VkParameters.fft == FFTEnum::C2C) - { - // For C2C computation we can do everything in the in-place-computation buffer. - inputGPUBuffer = GPUBuffer; - outputGPUBuffer = GPUBuffer; - } - else - { - if (m_VkParameters.I == DirectionEnum::FORWARD) + if (m_VkParameters.fft == FFTEnum::C2C) + { + // For C2C computation we can do everything in the in-place-computation buffer. + m_InputGPUBuffer = m_GPUBuffer; + m_OutputGPUBuffer = m_GPUBuffer; + } + else if (m_VkParameters.I == DirectionEnum::FORWARD) { - outputGPUBuffer = GPUBuffer; + m_OutputGPUBuffer = m_GPUBuffer; // Either R2FullH or R2HalfH. For forward computation, we have a smaller input buffer. const uint64_t inputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.inputBufferSize }; - resCu = cudaMalloc((void **)&inputGPUBuffer, inputBufferBytes); + resCu = cudaMalloc((void **)&m_InputGPUBuffer, inputBufferBytes); if (resCu != cudaSuccess) { std::cerr << __FILE__ "(" << __LINE__ << "): cudaMalloc returned " << resCu << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; } - - m_VkFFTConfiguration.inputBuffer = reinterpret_cast(&inputGPUBuffer); + m_VkFFTConfiguration.inputBuffer = reinterpret_cast(&m_InputGPUBuffer); } else { - inputGPUBuffer = GPUBuffer; + m_InputGPUBuffer = m_GPUBuffer; // Either R2FullH or R2HalfH. For inverse computation, we have a smaller output buffer. - uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; - resCu = cudaMalloc((void **)&outputGPUBuffer, outputBufferBytes); + const uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; + resCu = cudaMalloc((void **)&m_OutputGPUBuffer, outputBufferBytes); if (resCu != cudaSuccess) { std::cerr << __FILE__ "(" << __LINE__ << "): cudaMalloc returned " << resCu << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; } - m_VkFFTConfiguration.outputBuffer = reinterpret_cast(&outputGPUBuffer); + m_VkFFTConfiguration.outputBuffer = reinterpret_cast(&m_OutputGPUBuffer); } } - // Copy input from CPU to GPU - resCu = - cudaMemcpy(inputGPUBuffer, m_VkParameters.inputCPUBuffer, m_VkParameters.inputBufferBytes, cudaMemcpyHostToDevice); + // Copy input from CPU to GPU (per call) + resCu = cudaMemcpy( + m_InputGPUBuffer, m_VkParameters.inputCPUBuffer, m_VkParameters.inputBufferBytes, cudaMemcpyHostToDevice); if (resCu != cudaSuccess) { std::cerr << __FILE__ "(" << __LINE__ << "): cudaMemcpy returned " << resCu << std::endl; @@ -484,69 +497,68 @@ VkCommon::PerformFFT() #elif (VKFFT_BACKEND == OPENCL) cl_int resCL{ CL_SUCCESS }; - // Configure the buffers. Some of these three pointers will be nullptr or be duplicates of each - // other, so don't release all of them at the end. All re-striding of data (for R2HalfH or R2FullH, - // regardless of forward vs. inverse) is done by VkFFT between the two GPU buffers it uses. - cl_mem inputGPUBuffer{ nullptr }; // Copy from CPU input buffer to this GPU buffer - cl_mem GPUBuffer{ nullptr }; // GPU buffer where main computation occurs - cl_mem outputGPUBuffer{ nullptr }; // Copy from this GPU buffer to CPU output buffer - if (m_VkParameters.fft == FFTEnum::C2C) - { - // For C2C computation we can do everything in the in-place-computation buffer. - const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; - GPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, bufferBytes, nullptr, &resCL); - inputGPUBuffer = GPUBuffer; - outputGPUBuffer = GPUBuffer; - if (resCL != CL_SUCCESS) - { - std::cerr << __FILE__ "(" << __LINE__ << "): clCreateBuffer returned " << resCL << std::endl; - return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - } - m_VkFFTConfiguration.buffer = &GPUBuffer; - } - else + if (!m_PlanConfigured) { - // Either R2HalfH or R2FullH computation. Either forward or inverse. - const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; - GPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, bufferBytes, nullptr, &resCL); - if (resCL != CL_SUCCESS) + // Configure the persistent buffers for this shape. Some of m_*GPUBuffer alias each other + // (C2C, or the in-place direction), which the distinct-buffer free in ReleaseBackend handles. + if (m_VkParameters.fft == FFTEnum::C2C) { - std::cerr << __FILE__ "(" << __LINE__ << "): clCreateBuffer returned " << resCL << std::endl; - return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - } - m_VkFFTConfiguration.buffer = &GPUBuffer; - - if (m_VkParameters.I == DirectionEnum::FORWARD) - { - // Either R2FullH or R2HalfH. For forward computation, we have a smaller input buffer. - const uint64_t inputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.inputBufferSize }; - inputGPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, inputBufferBytes, nullptr, &resCL); - outputGPUBuffer = GPUBuffer; + // For C2C computation we can do everything in the in-place-computation buffer. + const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; + m_GPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, bufferBytes, nullptr, &resCL); + m_InputGPUBuffer = m_GPUBuffer; + m_OutputGPUBuffer = m_GPUBuffer; if (resCL != CL_SUCCESS) { std::cerr << __FILE__ "(" << __LINE__ << "): clCreateBuffer returned " << resCL << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; } - m_VkFFTConfiguration.inputBuffer = &inputGPUBuffer; + m_VkFFTConfiguration.buffer = &m_GPUBuffer; } else { - // Either R2FullH or R2HalfH. For inverse computation, we have a smaller output buffer. - uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; - inputGPUBuffer = GPUBuffer; - outputGPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, outputBufferBytes, nullptr, &resCL); + // Either R2HalfH or R2FullH computation. Either forward or inverse. + const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; + m_GPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, bufferBytes, nullptr, &resCL); if (resCL != CL_SUCCESS) { std::cerr << __FILE__ "(" << __LINE__ << "): clCreateBuffer returned " << resCL << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; } - m_VkFFTConfiguration.outputBuffer = &outputGPUBuffer; + m_VkFFTConfiguration.buffer = &m_GPUBuffer; + + if (m_VkParameters.I == DirectionEnum::FORWARD) + { + // Either R2FullH or R2HalfH. For forward computation, we have a smaller input buffer. + const uint64_t inputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.inputBufferSize }; + m_InputGPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, inputBufferBytes, nullptr, &resCL); + m_OutputGPUBuffer = m_GPUBuffer; + if (resCL != CL_SUCCESS) + { + std::cerr << __FILE__ "(" << __LINE__ << "): clCreateBuffer returned " << resCL << std::endl; + return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; + } + m_VkFFTConfiguration.inputBuffer = &m_InputGPUBuffer; + } + else + { + // Either R2FullH or R2HalfH. For inverse computation, we have a smaller output buffer. + const uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; + m_InputGPUBuffer = m_GPUBuffer; + m_OutputGPUBuffer = clCreateBuffer(m_VkGPU.context, CL_MEM_READ_WRITE, outputBufferBytes, nullptr, &resCL); + if (resCL != CL_SUCCESS) + { + std::cerr << __FILE__ "(" << __LINE__ << "): clCreateBuffer returned " << resCL << std::endl; + return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; + } + m_VkFFTConfiguration.outputBuffer = &m_OutputGPUBuffer; + } } } - // Copy input from CPU to GPU + // Copy input from CPU to GPU (per call) resCL = clEnqueueWriteBuffer(m_VkGPU.commandQueue, - inputGPUBuffer, + m_InputGPUBuffer, CL_TRUE, 0, m_VkParameters.inputBufferBytes, @@ -560,47 +572,48 @@ VkCommon::PerformFFT() return VkFFTResult{ VKFFT_ERROR_FAILED_TO_COPY }; } #elif (VKFFT_BACKEND == LEVEL_ZERO) - ze_result_t resZE{ ZE_RESULT_SUCCESS }; - void * inputGPUBuffer{ nullptr }; - void * GPUBuffer{ nullptr }; - void * outputGPUBuffer{ nullptr }; - ze_device_mem_alloc_desc_t deviceMemDesc{}; - deviceMemDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC; - - const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; - resZE = - zeMemAllocDevice(m_VkGPU.context, &deviceMemDesc, bufferBytes, m_VkParameters.PSize, m_VkGPU.device, &GPUBuffer); - if (resZE != ZE_RESULT_SUCCESS) - return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - m_VkFFTConfiguration.buffer = &GPUBuffer; + ze_result_t resZE{ ZE_RESULT_SUCCESS }; - if (m_VkParameters.fft == FFTEnum::C2C) - { - inputGPUBuffer = GPUBuffer; - outputGPUBuffer = GPUBuffer; - } - else if (m_VkParameters.I == DirectionEnum::FORWARD) - { - const uint64_t inputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.inputBufferSize }; - resZE = zeMemAllocDevice( - m_VkGPU.context, &deviceMemDesc, inputBufferBytes, m_VkParameters.PSize, m_VkGPU.device, &inputGPUBuffer); - if (resZE != ZE_RESULT_SUCCESS) - return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - outputGPUBuffer = GPUBuffer; - m_VkFFTConfiguration.inputBuffer = &inputGPUBuffer; - } - else + if (!m_PlanConfigured) { - const uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; + ze_device_mem_alloc_desc_t deviceMemDesc{}; + deviceMemDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC; + + const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; resZE = zeMemAllocDevice( - m_VkGPU.context, &deviceMemDesc, outputBufferBytes, m_VkParameters.PSize, m_VkGPU.device, &outputGPUBuffer); + m_VkGPU.context, &deviceMemDesc, bufferBytes, m_VkParameters.PSize, m_VkGPU.device, &m_GPUBuffer); if (resZE != ZE_RESULT_SUCCESS) return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - inputGPUBuffer = GPUBuffer; - m_VkFFTConfiguration.outputBuffer = &outputGPUBuffer; + m_VkFFTConfiguration.buffer = &m_GPUBuffer; + + if (m_VkParameters.fft == FFTEnum::C2C) + { + m_InputGPUBuffer = m_GPUBuffer; + m_OutputGPUBuffer = m_GPUBuffer; + } + else if (m_VkParameters.I == DirectionEnum::FORWARD) + { + const uint64_t inputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.inputBufferSize }; + resZE = zeMemAllocDevice( + m_VkGPU.context, &deviceMemDesc, inputBufferBytes, m_VkParameters.PSize, m_VkGPU.device, &m_InputGPUBuffer); + if (resZE != ZE_RESULT_SUCCESS) + return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; + m_OutputGPUBuffer = m_GPUBuffer; + m_VkFFTConfiguration.inputBuffer = &m_InputGPUBuffer; + } + else + { + const uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; + resZE = zeMemAllocDevice( + m_VkGPU.context, &deviceMemDesc, outputBufferBytes, m_VkParameters.PSize, m_VkGPU.device, &m_OutputGPUBuffer); + if (resZE != ZE_RESULT_SUCCESS) + return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; + m_InputGPUBuffer = m_GPUBuffer; + m_VkFFTConfiguration.outputBuffer = &m_OutputGPUBuffer; + } } - // Host -> device copy via an immediate command list on the compute/copy queue group. + // Host -> device copy via an immediate command list on the compute/copy queue group (per call). { ze_command_queue_desc_t copyQueueDesc{}; copyQueueDesc.stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_DESC; @@ -612,7 +625,7 @@ VkCommon::PerformFFT() if (resZE != ZE_RESULT_SUCCESS) return VkFFTResult{ VKFFT_ERROR_FAILED_TO_CREATE_COMMAND_LIST }; resZE = zeCommandListAppendMemoryCopy(copyCommandList, - inputGPUBuffer, + m_InputGPUBuffer, m_VkParameters.inputCPUBuffer, m_VkParameters.inputBufferBytes, nullptr, @@ -628,57 +641,63 @@ VkCommon::PerformFFT() #elif (VKFFT_BACKEND == METAL) // Metal shared-storage buffers are CPU-visible on Apple unified-memory systems, // so host<->device transfers reduce to memcpy into/out of MTL::Buffer::contents(). - MTL::Buffer * inputGPUBuffer{ nullptr }; - MTL::Buffer * GPUBuffer{ nullptr }; - MTL::Buffer * outputGPUBuffer{ nullptr }; - const auto storageMode = MTL::ResourceStorageModeShared; - if (m_VkParameters.fft == FFTEnum::C2C) - { - const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; - GPUBuffer = m_VkGPU.device->newBuffer(bufferBytes, storageMode); - if (GPUBuffer == nullptr) - return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - inputGPUBuffer = GPUBuffer; - outputGPUBuffer = GPUBuffer; - m_VkFFTConfiguration.buffer = &GPUBuffer; - } - else + const auto storageMode = MTL::ResourceStorageModeShared; + if (!m_PlanConfigured) { - const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; - GPUBuffer = m_VkGPU.device->newBuffer(bufferBytes, storageMode); - if (GPUBuffer == nullptr) - return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - m_VkFFTConfiguration.buffer = &GPUBuffer; - - if (m_VkParameters.I == DirectionEnum::FORWARD) + if (m_VkParameters.fft == FFTEnum::C2C) { - const uint64_t inputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.inputBufferSize }; - inputGPUBuffer = m_VkGPU.device->newBuffer(inputBufferBytes, storageMode); - if (inputGPUBuffer == nullptr) + const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; + m_GPUBuffer = m_VkGPU.device->newBuffer(bufferBytes, storageMode); + if (m_GPUBuffer == nullptr) return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - outputGPUBuffer = GPUBuffer; - m_VkFFTConfiguration.inputBuffer = &inputGPUBuffer; + m_InputGPUBuffer = m_GPUBuffer; + m_OutputGPUBuffer = m_GPUBuffer; + m_VkFFTConfiguration.buffer = &m_GPUBuffer; } else { - const uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; - inputGPUBuffer = GPUBuffer; - outputGPUBuffer = m_VkGPU.device->newBuffer(outputBufferBytes, storageMode); - if (outputGPUBuffer == nullptr) + const uint64_t bufferBytes{ 2UL * m_VkParameters.PSize * *m_VkFFTConfiguration.bufferSize }; + m_GPUBuffer = m_VkGPU.device->newBuffer(bufferBytes, storageMode); + if (m_GPUBuffer == nullptr) return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; - m_VkFFTConfiguration.outputBuffer = &outputGPUBuffer; + m_VkFFTConfiguration.buffer = &m_GPUBuffer; + + if (m_VkParameters.I == DirectionEnum::FORWARD) + { + const uint64_t inputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.inputBufferSize }; + m_InputGPUBuffer = m_VkGPU.device->newBuffer(inputBufferBytes, storageMode); + if (m_InputGPUBuffer == nullptr) + return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; + m_OutputGPUBuffer = m_GPUBuffer; + m_VkFFTConfiguration.inputBuffer = &m_InputGPUBuffer; + } + else + { + const uint64_t outputBufferBytes{ 1UL * m_VkParameters.PSize * *m_VkFFTConfiguration.outputBufferSize }; + m_InputGPUBuffer = m_GPUBuffer; + m_OutputGPUBuffer = m_VkGPU.device->newBuffer(outputBufferBytes, storageMode); + if (m_OutputGPUBuffer == nullptr) + return VkFFTResult{ VKFFT_ERROR_FAILED_TO_ALLOCATE }; + m_VkFFTConfiguration.outputBuffer = &m_OutputGPUBuffer; + } } } - std::memcpy(inputGPUBuffer->contents(), m_VkParameters.inputCPUBuffer, m_VkParameters.inputBufferBytes); + std::memcpy(m_InputGPUBuffer->contents(), m_VkParameters.inputCPUBuffer, m_VkParameters.inputBufferBytes); #endif - // Initialize applications. This function loads shaders, creates pipeline and configures FFT based on configuration - // file. No buffer allocations inside VkFFT library. - VkFFTApplication app{}; - resFFT = initializeVkFFT(&app, m_VkFFTConfiguration); - if (resFFT != VKFFT_SUCCESS) - return resFFT; + // Initialize the application once per plan. This loads shaders, creates the pipeline, and + // compiles the FFT kernels (nvrtc for CUDA, clBuildProgram for OpenCL); it allocates no + // buffers. The compiled plan and the persistent buffers above are reused across same-shape + // calls -- buffers are bound per call through VkFFTLaunchParams below. + if (!m_PlanConfigured) + { + m_VkFFTApplication = new VkFFTApplication{}; + resFFT = initializeVkFFT(m_VkFFTApplication, m_VkFFTConfiguration); + if (resFFT != VKFFT_SUCCESS) + return resFFT; + m_PlanConfigured = true; + } // Submit FFT or iFFT. VkFFTLaunchParams launchParams{}; @@ -705,7 +724,7 @@ VkCommon::PerformFFT() launchParams.commandEncoder = metalEncoder; #endif - resFFT = VkFFTAppend(&app, m_VkParameters.I == DirectionEnum::INVERSE ? 1 : -1, &launchParams); + resFFT = VkFFTAppend(m_VkFFTApplication, m_VkParameters.I == DirectionEnum::INVERSE ? 1 : -1, &launchParams); if (resFFT != VKFFT_SUCCESS) return resFFT; @@ -717,22 +736,15 @@ VkCommon::PerformFFT() return VkFFTResult{ VKFFT_ERROR_FAILED_TO_SYNCHRONIZE }; } - // Copy result from GPU to CPU + // Copy result from GPU to CPU (per call); persistent buffers are freed in ReleaseBackend(). resCu = cudaMemcpy( - m_VkParameters.outputCPUBuffer, outputGPUBuffer, m_VkParameters.outputBufferBytes, cudaMemcpyDeviceToHost); + m_VkParameters.outputCPUBuffer, m_OutputGPUBuffer, m_VkParameters.outputBufferBytes, cudaMemcpyDeviceToHost); if (resCu != cudaSuccess) { std::cerr << __FILE__ "(" << __LINE__ << "): cudaMemcpy returned " << resCu << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_COPY }; } - // Release mem buffers - cudaFree(inputGPUBuffer); - if (m_VkParameters.fft != FFTEnum::C2C) - { - cudaFree(outputGPUBuffer); - } - #elif (VKFFT_BACKEND == OPENCL) resCL = clFinish(m_VkGPU.commandQueue); if (resCL != CL_SUCCESS) @@ -741,9 +753,9 @@ VkCommon::PerformFFT() return VkFFTResult{ VKFFT_ERROR_FAILED_TO_SYNCHRONIZE }; } - // Copy result from GPU to CPU + // Copy result from GPU to CPU (per call); persistent buffers are freed in ReleaseBackend(). resCL = clEnqueueReadBuffer(m_VkGPU.commandQueue, - outputGPUBuffer, + m_OutputGPUBuffer, CL_TRUE, 0, m_VkParameters.outputBufferBytes, @@ -756,13 +768,6 @@ VkCommon::PerformFFT() std::cerr << __FILE__ "(" << __LINE__ << "): clEnqueueReadBuffer returned " << resCL << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_COPY }; } - - clReleaseMemObject(inputGPUBuffer); - if (m_VkParameters.fft != FFTEnum::C2C) - { - // Release other buffer too - clReleaseMemObject(outputGPUBuffer); - } #elif (VKFFT_BACKEND == LEVEL_ZERO) resZE = zeCommandListClose(launchCommandList); if (resZE != ZE_RESULT_SUCCESS) @@ -788,7 +793,7 @@ VkCommon::PerformFFT() return VkFFTResult{ VKFFT_ERROR_FAILED_TO_CREATE_COMMAND_LIST }; resZE = zeCommandListAppendMemoryCopy(copyCommandList, m_VkParameters.outputCPUBuffer, - outputGPUBuffer, + m_OutputGPUBuffer, m_VkParameters.outputBufferBytes, nullptr, 0, @@ -800,32 +805,14 @@ VkCommon::PerformFFT() return VkFFTResult{ VKFFT_ERROR_FAILED_TO_SYNCHRONIZE }; zeCommandListDestroy(copyCommandList); } - - // GPUBuffer aliases input or output in the C2C / inverse-R2H cases; free it once. - zeMemFree(m_VkGPU.context, GPUBuffer); - if (m_VkParameters.fft != FFTEnum::C2C) - { - if (m_VkParameters.I == DirectionEnum::FORWARD) - zeMemFree(m_VkGPU.context, inputGPUBuffer); - else - zeMemFree(m_VkGPU.context, outputGPUBuffer); - } + // Persistent buffers are freed in ReleaseBackend(). #elif (VKFFT_BACKEND == METAL) metalEncoder->endEncoding(); metalCommandBuffer->commit(); metalCommandBuffer->waitUntilCompleted(); - std::memcpy(m_VkParameters.outputCPUBuffer, outputGPUBuffer->contents(), m_VkParameters.outputBufferBytes); - - // The C2C in-place case aliases input/output to GPUBuffer; release once. - GPUBuffer->release(); - if (m_VkParameters.fft != FFTEnum::C2C) - { - if (m_VkParameters.I == DirectionEnum::FORWARD) - inputGPUBuffer->release(); - else - outputGPUBuffer->release(); - } + std::memcpy(m_VkParameters.outputCPUBuffer, m_OutputGPUBuffer->contents(), m_VkParameters.outputBufferBytes); + // Persistent buffers are released in ReleaseBackend(). #endif if (m_VkParameters.fft == FFTEnum::R2FullH && m_VkParameters.I == DirectionEnum::FORWARD) @@ -873,7 +860,8 @@ VkCommon::PerformFFT() break; } // end switch (m_VkParameters.P) } // end if(m_VkParameters.fft == R2FullH && m_VkParameters.I == DirectionEnum::FORWARD) - deleteVkFFT(&app); + // The plan (m_VkFFTApplication) and the GPU buffers persist for reuse and are released in + // ReleaseBackend() on the next device/shape change or at destruction. return resFFT; } @@ -883,15 +871,51 @@ VkCommon::ReleaseBackend() { VkFFTResult resFFT{ VKFFT_SUCCESS }; - // Return to launchVkFFT code #if (VKFFT_BACKEND == CUDA) + // Make this instance's context current so deleteVkFFT and the buffer frees below operate on + // the context that owns these GPU resources, not whichever filter configured last. + if (m_VkGPU.context) + cuCtxSetCurrent(m_VkGPU.context); +#endif + + // Release the cached plan and persistent buffers before the context is destroyed; their GPU + // resources belong to it. Distinct, non-aliased buffers are each freed exactly once. + if (m_PlanConfigured) + { + deleteVkFFT(m_VkFFTApplication); + delete m_VkFFTApplication; + m_VkFFTApplication = nullptr; + m_PlanConfigured = false; + } + +#if (VKFFT_BACKEND == CUDA) + if (m_GPUBuffer) + cudaFree(m_GPUBuffer); + if (m_InputGPUBuffer && m_InputGPUBuffer != m_GPUBuffer) + cudaFree(m_InputGPUBuffer); + if (m_OutputGPUBuffer && m_OutputGPUBuffer != m_GPUBuffer && m_OutputGPUBuffer != m_InputGPUBuffer) + cudaFree(m_OutputGPUBuffer); + m_GPUBuffer = nullptr; + m_InputGPUBuffer = nullptr; + m_OutputGPUBuffer = nullptr; if (m_VkGPU.context) { cuCtxDestroy(m_VkGPU.context); + m_VkGPU.context = 0; } #elif (VKFFT_BACKEND == OPENCL) cl_int resCL{ CL_SUCCESS }; + if (m_GPUBuffer) + clReleaseMemObject(m_GPUBuffer); + if (m_InputGPUBuffer && m_InputGPUBuffer != m_GPUBuffer) + clReleaseMemObject(m_InputGPUBuffer); + if (m_OutputGPUBuffer && m_OutputGPUBuffer != m_GPUBuffer && m_OutputGPUBuffer != m_InputGPUBuffer) + clReleaseMemObject(m_OutputGPUBuffer); + m_GPUBuffer = nullptr; + m_InputGPUBuffer = nullptr; + m_OutputGPUBuffer = nullptr; + if (m_VkGPU.commandQueue) { resCL = clReleaseCommandQueue(m_VkGPU.commandQueue); @@ -900,6 +924,7 @@ VkCommon::ReleaseBackend() std::cerr << __FILE__ "(" << __LINE__ << "): clReleaseCommandQueue returned " << resCL << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_RELEASE_COMMAND_QUEUE }; } + m_VkGPU.commandQueue = 0; } if (m_VkGPU.context) @@ -910,8 +935,22 @@ VkCommon::ReleaseBackend() std::cerr << __FILE__ "(" << __LINE__ << "): clReleaseContext returned " << resCL << std::endl; return VkFFTResult{ VKFFT_ERROR_FAILED_TO_RELEASE_COMMAND_QUEUE }; } + m_VkGPU.context = 0; } #elif (VKFFT_BACKEND == LEVEL_ZERO) + if (m_VkGPU.context) + { + if (m_GPUBuffer) + zeMemFree(m_VkGPU.context, m_GPUBuffer); + if (m_InputGPUBuffer && m_InputGPUBuffer != m_GPUBuffer) + zeMemFree(m_VkGPU.context, m_InputGPUBuffer); + if (m_OutputGPUBuffer && m_OutputGPUBuffer != m_GPUBuffer && m_OutputGPUBuffer != m_InputGPUBuffer) + zeMemFree(m_VkGPU.context, m_OutputGPUBuffer); + } + m_GPUBuffer = nullptr; + m_InputGPUBuffer = nullptr; + m_OutputGPUBuffer = nullptr; + if (m_VkGPU.commandQueue) { zeCommandQueueDestroy(m_VkGPU.commandQueue); @@ -923,6 +962,16 @@ VkCommon::ReleaseBackend() m_VkGPU.context = nullptr; } #elif (VKFFT_BACKEND == METAL) + if (m_GPUBuffer) + m_GPUBuffer->release(); + if (m_InputGPUBuffer && m_InputGPUBuffer != m_GPUBuffer) + m_InputGPUBuffer->release(); + if (m_OutputGPUBuffer && m_OutputGPUBuffer != m_GPUBuffer && m_OutputGPUBuffer != m_InputGPUBuffer) + m_OutputGPUBuffer->release(); + m_GPUBuffer = nullptr; + m_InputGPUBuffer = nullptr; + m_OutputGPUBuffer = nullptr; + if (m_VkGPU.queue) { m_VkGPU.queue->release(); diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 32ac8f9..0635a95 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -7,6 +7,7 @@ set( itkVkComplexToComplex1DFFTImageFilterSizesTest.cxx itkVkDiscreteGaussianImageFilterTest.cxx itkVkFFTImageFilterFactoryTest.cxx + itkVkFFTRepeatedTransformLeakTest.cxx itkVkForwardInverseFFTImageFilterTest.cxx itkVkForwardInverse1DFFTImageFilterTest.cxx itkVkForward1DFFTImageFilterBaselineTest.cxx @@ -176,6 +177,13 @@ itk_add_test(NAME itkVkForwardInverseFFTImageFilterTestDouble ) _vkfft_disable_on_unsupported_fp64(itkVkForwardInverseFFTImageFilterTestDouble) +# Regression guard for the per-Update GPU context leak: reuse one filter for +# many transforms (precision, size^3, iterations). Before context caching this +# aborted with a VkFFT out-of-memory error after a few dozen iterations. +itk_add_test(NAME itkVkFFTRepeatedTransformLeakTestFloat + COMMAND VkFFTBackendTestDriver itkVkFFTRepeatedTransformLeakTest float 64 200 +) + # pocl (CPU OpenCL) computes VkFFT's size-19 Bluestein inverse incorrectly, so # the hosted pocl-backed leg runs the round-trip tests capped at size 16 # (radix-2/3/5/7 plus Bluestein primes 11/13). Real GPUs run the full sweep above. diff --git a/test/itkVkFFTRepeatedTransformLeakTest.cxx b/test/itkVkFFTRepeatedTransformLeakTest.cxx new file mode 100644 index 0000000..2de98d8 --- /dev/null +++ b/test/itkVkFFTRepeatedTransformLeakTest.cxx @@ -0,0 +1,109 @@ +/*========================================================================= + * + * Copyright NumFOCUS + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0.txt + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + *=========================================================================*/ + +// Regression test for the per-Update GPU resource leak. +// +// A single reused Vk FFT filter driven for many transforms (the pattern an +// iterative registration uses) previously created and leaked one GPU context +// per Update(), exhausting device memory and aborting with a VkFFT +// out-of-memory error after a few dozen iterations. With the context cached +// across calls, device memory stays flat and the loop runs to completion. + +#include "itkVkRealToHalfHermitianForwardFFTImageFilter.h" +#include "itkVkHalfHermitianToRealInverseFFTImageFilter.h" + +#include "itkImageRegionIterator.h" +#include "itkTestingMacros.h" + +#include +#include +#include + +template +int +runRepeatedTransformLeakTest(unsigned int size, unsigned int iterations) +{ + constexpr unsigned int Dimension{ 3 }; + using RealImageType = itk::Image; + using ComplexImageType = itk::Image, Dimension>; + + typename RealImageType::SizeType regionSize; + regionSize.Fill(size); + auto input = RealImageType::New(); + input->SetRegions(regionSize); + input->Allocate(); + input->FillBuffer(PrecisionType{ 1 }); + + // A single forward+inverse pair reused for every iteration, exactly as a + // registration loop reuses its smoothing filters. + using ForwardFilterType = itk::VkRealToHalfHermitianForwardFFTImageFilter; + using InverseFilterType = itk::VkHalfHermitianToRealInverseFFTImageFilter; + auto forwardFilter = ForwardFilterType::New(); + auto inverseFilter = InverseFilterType::New(); + forwardFilter->SetDeviceID(0); + inverseFilter->SetDeviceID(0); + forwardFilter->SetInput(input); + inverseFilter->SetInput(forwardFilter->GetOutput()); + + for (unsigned int i{ 0 }; i < iterations; ++i) + { + // Fresh input data each iteration forces a real recomputation (new CPU + // buffer contents), reproducing the per-call pattern that leaked. + input->FillBuffer(static_cast((i % 7) + 1)); + input->Modified(); + ITK_TRY_EXPECT_NO_EXCEPTION(inverseFilter->UpdateLargestPossibleRegion()); + } + + // Round-trip sanity: inverse(forward(constant)) is the same constant. + inverseFilter->GetOutput()->Update(); + const auto lastValue = static_cast(((iterations - 1) % 7) + 1); + itk::ImageRegionIterator it(inverseFilter->GetOutput(), + inverseFilter->GetOutput()->GetBufferedRegion()); + const PrecisionType tolerance{ static_cast(1e-3) }; + for (it.GoToBegin(); !it.IsAtEnd(); ++it) + { + if (std::abs(it.Get() - lastValue) > tolerance) + { + std::cerr << "Round-trip mismatch: got " << it.Get() << " expected " << lastValue << std::endl; + return EXIT_FAILURE; + } + } + + std::cout << "Completed " << iterations << " forward+inverse transforms at " << size << "^3 without GPU resource " + << "exhaustion." << std::endl; + return EXIT_SUCCESS; +} + +int +itkVkFFTRepeatedTransformLeakTest(int argc, char * argv[]) +{ + const std::string precision{ (argc > 1) ? argv[1] : "float" }; + const unsigned int size{ (argc > 2) ? static_cast(std::stoul(argv[2])) : 64 }; + const unsigned int iterations{ (argc > 3) ? static_cast(std::stoul(argv[3])) : 200 }; + + if (precision == "double") + { + return runRepeatedTransformLeakTest(size, iterations); + } + if (precision == "float") + { + return runRepeatedTransformLeakTest(size, iterations); + } + std::cerr << "Unknown precision '" << precision << "'. Expected 'float' or 'double'." << std::endl; + return EXIT_FAILURE; +}