Merge pull request #17 from tbirdso/1d-fft

ENH: Add accelerated 1D FFT classes

.github/workflows/test-gpu.yml

@@ -12,7 +12,7 @@ jobs:
             opencl-headers-git-tag: "v2021.04.29"
             opencl-version: 120
             vkfft-backend: 3
-            itk-git-tag: "c693dbc830859ea10d573301f075e121a4c6c581"
+            itk-git-tag: "b1a0a91614dc091f9a17c1e4727095c80e9c05ea"
             cmake-build-type: "MinSizeRel"
             platform-name: "ubuntu-nvidia-gpu"
             os: ubuntu-20.04

include/itkVkCommon.h

@@ -63,9 +63,13 @@ class VkFFTBackend_EXPORT VkCommon
 
   struct VkParameters
   {
-    uint64_t      X = 0;                    // size of fastest varying dimension
-    uint64_t      Y = 1;                    // size of second-fastest varying dimension, if any, otherwise 1.
-    uint64_t      Z = 1;                    // size of third-fastest varying dimension, if any, otherwise 1.
+    uint64_t X = 0; // size of fastest varying dimension
+    uint64_t Y = 1; // size of second-fastest varying dimension, if any, otherwise 1.
+    uint64_t Z = 1; // size of third-fastest varying dimension, if any, otherwise 1.
+    uint64_t omitDimension[3] = { 0,
+                                  0,
+                                  0 }; // disable FFT for this dimension (0 - FFT enabled, 1 - FFT disabled). Default 0.
+                                       // Doesn't work for R2C dimension 0 for now. Doesn't work with convolutions.
     PrecisionEnum P = PrecisionEnum::FLOAT; // type for real numbers
     uint64_t      B = 1;                    // Number of batches -- always 1
     uint64_t      N = 1;                    // Number of redundant iterations, for benchmarking -- always 1.

include/itkVkComplexToComplex1DFFTImageFilter.h

@@ -0,0 +1,108 @@
+/*=========================================================================
+ *
+ *  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
+ *
+ *         http://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.
+ *
+ *=========================================================================*/
+
+#ifndef itkVkComplexToComplex1DFFTImageFilter_h
+#define itkVkComplexToComplex1DFFTImageFilter_h
+
+#include "itkVkCommon.h"
+#include "itkComplexToComplex1DFFTImageFilter.h"
+
+namespace itk
+{
+
+/**
+ *\class VkComplexToComplex1DFFTImageFilter
+ *
+ *  \brief Implements an API to enable the 1D Fourier transform or the inverse
+ *  Fourier transform of images with complex valued voxels to be computed using
+ *  the VkFFT library.
+ *
+ * This filter is multithreaded and by default supports input images with sizes which are
+ * divisible by primes up to 13. 
+ * 
+ * Execution on input images with sizes divisible by primes greater than 17 may succeed 
+ * with a fallback on Bluestein's algorithm per VkFFT with a cost to performance and output precision.
+ *
+ * \ingroup FourierTransform
+ * \ingroup MultiThreaded
+ * \ingroup ITKFFT
+ * \ingroup VkFFTBackend
+ */
+template <typename TImage>
+class VkComplexToComplex1DFFTImageFilter : public ComplexToComplex1DFFTImageFilter<TImage>
+{
+public:
+  ITK_DISALLOW_COPY_AND_MOVE(VkComplexToComplex1DFFTImageFilter);
+
+  using InputImageType = TImage;
+  using OutputImageType = TImage;
+  static_assert(std::is_same<typename TImage::PixelType, std::complex<float>>::value ||
+                  std::is_same<typename TImage::PixelType, std::complex<double>>::value,
+                "Unsupported pixel type");
+  static_assert(TImage::ImageDimension >= 1 && TImage::ImageDimension <= 3, "Unsupported image dimension");
+
+  /** Standard class type aliases. */
+  using Self = VkComplexToComplex1DFFTImageFilter;
+  using Superclass = ComplexToComplex1DFFTImageFilter<TImage>;
+  using Pointer = SmartPointer<Self>;
+  using ConstPointer = SmartPointer<const Self>;
+  using InputPixelType = typename InputImageType::PixelType;
+  using OutputPixelType = typename OutputImageType::PixelType;
+  using ComplexType = InputPixelType;
+  using RealType = typename ComplexType::value_type;
+  using SizeType = typename InputImageType::SizeType;
+  using SizeValueType = typename InputImageType::SizeValueType;
+  using OutputImageRegionType = typename OutputImageType::RegionType;
+
+  /** Method for creation through the object factory. */
+  itkNewMacro(Self);
+
+  /** Run-time type information (and related methods). */
+  itkTypeMacro(VkComplexToComplex1DFFTImageFilter, ComplexToComplex1DFFTImageFilter);
+
+  static constexpr unsigned int ImageDimension = InputImageType::ImageDimension;
+
+  itkGetMacro(DeviceID, uint64_t);
+  itkSetMacro(DeviceID, uint64_t);
+
+  SizeValueType
+  GetSizeGreatestPrimeFactor() const override;
+
+protected:
+  VkComplexToComplex1DFFTImageFilter() = default;
+  ~VkComplexToComplex1DFFTImageFilter() override = default;
+
+  void
+  GenerateData() override;
+
+  void
+  PrintSelf(std::ostream & os, Indent indent) const override;
+
+private:
+  uint64_t m_DeviceID{ 0UL };
+
+  VkCommon m_VkCommon{};
+};
+
+} // namespace itk
+
+#ifndef ITK_MANUAL_INSTANTIATION
+#  include "itkVkComplexToComplex1DFFTImageFilter.hxx"
+#endif
+
+#endif // itkVkComplexToComplex1DFFTImageFilter_h

include/itkVkComplexToComplex1DFFTImageFilter.hxx

@@ -0,0 +1,127 @@
+/*=========================================================================
+ *
+ *  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
+ *
+ *         http://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.
+ *
+ *=========================================================================*/
+
+#ifndef itkVkComplexToComplex1DFFTImageFilter_hxx
+#define itkVkComplexToComplex1DFFTImageFilter_hxx
+
+#include "itkVkComplexToComplex1DFFTImageFilter.h"
+#include "vkFFT.h"
+#include "itkImageRegionIterator.h"
+#include "itkIndent.h"
+#include "itkMetaDataObject.h"
+#include "itkProgressReporter.h"
+
+namespace itk
+{
+template <typename TImage>
+void
+VkComplexToComplex1DFFTImageFilter<TImage>::GenerateData()
+{
+  // get pointers to the input and output
+  const InputImageType * const input{ this->GetInput() };
+  OutputImageType * const      output{ this->GetOutput() };
+
+  if (!input || !output)
+  {
+    return;
+  }
+
+  // we don't have a nice progress to report, but at least this simple line
+  // reports the beginning and the end of the process
+  const ProgressReporter progress(this, 0, 1);
+
+  // allocate output buffer memory
+  output->SetBufferedRegion(output->GetRequestedRegion());
+  output->Allocate();
+
+  const SizeType & inputSize{ input->GetLargestPossibleRegion().GetSize() };
+
+  const InputPixelType * const inputCPUBuffer{ input->GetBufferPointer() };
+  OutputPixelType * const      outputCPUBuffer{ output->GetBufferPointer() };
+  itkAssertOrThrowMacro(inputCPUBuffer != nullptr, "No CPU input buffer");
+  itkAssertOrThrowMacro(outputCPUBuffer != nullptr, "No CPU output buffer");
+  const SizeValueType inBytes{ input->GetLargestPossibleRegion().GetNumberOfPixels() * sizeof(InputPixelType) };
+  const SizeValueType outBytes{ output->GetLargestPossibleRegion().GetNumberOfPixels() * sizeof(OutputPixelType) };
+  itkAssertOrThrowMacro(inBytes == outBytes, "CPU input and output buffers are of different sizes.");
+
+  // Mostly use defaults for VkCommon::VkGPU
+  typename VkCommon::VkGPU vkGPU;
+  vkGPU.device_id = m_DeviceID;
+
+  // Describe this filter in VkCommon::VkParameters
+  typename VkCommon::VkParameters vkParameters;
+  if (ImageDimension > 0)
+    vkParameters.X = inputSize[0];
+  if (ImageDimension > 1)
+    vkParameters.Y = inputSize[1];
+  if (ImageDimension > 2)
+    vkParameters.Z = inputSize[2];
+  if (std::is_same<RealType, float>::value)
+    vkParameters.P = VkCommon::PrecisionEnum::FLOAT;
+  else if (std::is_same<RealType, double>::value)
+    vkParameters.P = VkCommon::PrecisionEnum::DOUBLE;
+  else
+    itkAssertOrThrowMacro(false, "Unsupported type for real numbers.");
+  vkParameters.fft = VkCommon::FFTEnum::C2C;
+  vkParameters.PSize = sizeof(RealType);
+  vkParameters.I = this->GetTransformDirection() == Superclass::TransformDirectionType::INVERSE
+                     ? VkCommon::DirectionEnum::INVERSE
+                     : VkCommon::DirectionEnum::FORWARD;
+  vkParameters.normalized = vkParameters.I == VkCommon::DirectionEnum::INVERSE
+                              ? VkCommon::NormalizationEnum::NORMALIZED
+                              : VkCommon::NormalizationEnum::UNNORMALIZED;
+  for (size_t dim = 0; dim < ImageDimension; ++dim)
+  {
+    if (this->GetDirection() != dim)
+    {
+      vkParameters.omitDimension[dim] = 1; // omit dimensions other than in the given direction.
+    }
+  }
+
+  vkParameters.inputCPUBuffer = inputCPUBuffer;
+  vkParameters.inputBufferBytes = inBytes;
+  vkParameters.outputCPUBuffer = outputCPUBuffer;
+  vkParameters.outputBufferBytes = outBytes;
+
+  const VkFFTResult resFFT{ m_VkCommon.Run(vkGPU, vkParameters) };
+  if (resFFT != VKFFT_SUCCESS)
+  {
+    std::ostringstream mesg;
+    mesg << "VkFFT third-party library failed with error code " << resFFT << ".";
+    itkAssertOrThrowMacro(false, mesg.str());
+  }
+}
+
+template <typename TImage>
+void
+VkComplexToComplex1DFFTImageFilter<TImage>::PrintSelf(std::ostream & os, Indent indent) const
+{
+  Superclass::PrintSelf(os, indent);
+  os << indent << "DeviceID: " << m_DeviceID << std::endl;
+}
+
+template <typename TImage>
+typename VkComplexToComplex1DFFTImageFilter<TImage>::SizeValueType
+VkComplexToComplex1DFFTImageFilter<TImage>::GetSizeGreatestPrimeFactor() const
+{
+  return m_VkCommon.GetGreatestPrimeFactor();
+}
+
+} // end namespace itk
+
+#endif // _itkVkComplexToComplex1DFFTImageFilter_hxx

include/itkVkForward1DFFTImageFilter.h

@@ -0,0 +1,113 @@
+/*=========================================================================
+ *
+ *  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
+ *
+ *         http://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.
+ *
+ *=========================================================================*/
+#ifndef itkVkForward1DFFTImageFilter_h
+#define itkVkForward1DFFTImageFilter_h
+
+#include "itkForward1DFFTImageFilter.h"
+#include "itkVkCommon.h"
+
+namespace itk
+{
+/**
+ *\class VkForward1DFFTImageFilter
+ *
+ * \brief Vk-based forward 1D Fast Fourier Transform.
+ *
+ * This filter computes the forward Fourier transform in one dimension of an image. The
+ * implementation is based on the VkFFT library.
+ *
+ * This filter is multithreaded and by default supports input images with sizes which are
+ * divisible by primes up to 13.
+ * 
+ * Execution on input images with sizes divisible by primes greater than 17 may succeed 
+ * with a fallback on Bluestein's algorithm per VkFFT with a cost to performance and output precision.
+ *
+ * \ingroup FourierTransform
+ * \ingroup MultiThreaded
+ * \ingroup ITKFFT
+ * \ingroup VkFFTBackend
+ *
+ * \sa VkGlobalConfiguration
+ * \sa Forward1DFFTImageFilter
+ */
+template <typename TInputImage>
+class VkForward1DFFTImageFilter
+  : public Forward1DFFTImageFilter<TInputImage,
+                                 Image<std::complex<typename TInputImage::PixelType>, TInputImage::ImageDimension>>
+{
+public:
+  ITK_DISALLOW_COPY_AND_MOVE(VkForward1DFFTImageFilter);
+
+  using InputImageType = TInputImage;
+  using OutputImageType = Image<std::complex<typename TInputImage::PixelType>, TInputImage::ImageDimension>;
+  static_assert(std::is_same<typename TInputImage::PixelType, float>::value ||
+                  std::is_same<typename TInputImage::PixelType, double>::value,
+                "Unsupported pixel type");
+  static_assert(TInputImage::ImageDimension >= 1 && TInputImage::ImageDimension <= 3, "Unsupported image dimension");
+
+  /** Standard class type aliases. */
+  using Self = VkForward1DFFTImageFilter;
+  using Superclass = Forward1DFFTImageFilter<InputImageType, OutputImageType>;
+  using Pointer = SmartPointer<Self>;
+  using ConstPointer = SmartPointer<const Self>;
+
+  using InputPixelType = typename InputImageType::PixelType;
+  using OutputPixelType = typename OutputImageType::PixelType;
+  using ComplexType = OutputPixelType;
+  using RealType = typename ComplexType::value_type;
+  using SizeType = typename InputImageType::SizeType;
+  using SizeValueType = typename InputImageType::SizeValueType;
+  using OutputImageRegionType = typename OutputImageType::RegionType;
+
+  /** Method for creation through the object factory. */
+  itkNewMacro(Self);
+
+  /** Run-time type information (and related methods). */
+  itkTypeMacro(VkForward1DFFTImageFilter, Forward1DFFTImageFilter);
+
+  static constexpr unsigned int ImageDimension = InputImageType::ImageDimension;
+
+  itkGetMacro(DeviceID, uint64_t);
+  itkSetMacro(DeviceID, uint64_t);
+
+  SizeValueType
+  GetSizeGreatestPrimeFactor() const override;
+
+protected:
+  VkForward1DFFTImageFilter() = default;
+  ~VkForward1DFFTImageFilter() override = default;
+
+  void
+  GenerateData() override;
+
+  void
+  PrintSelf(std::ostream & os, Indent indent) const override;
+
+private:
+  uint64_t m_DeviceID{ 0UL };
+
+  VkCommon m_VkCommon{};
+};
+
+} // namespace itk
+
+#ifndef ITK_MANUAL_INSTANTIATION
+#  include "itkVkForward1DFFTImageFilter.hxx"
+#endif
+
+#endif // itkVkForward1DFFTImageFilter_h