Fixed bug with FFT size shorter than length of tensor (#270)

NVIDIA · Sep 8, 2022 · 4c5ddde · 4c5ddde
1 parent 8216b2a
commit 4c5ddde
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git a/include/matx/transforms/fft.h b/include/matx/transforms/fft.h
@@ -704,7 +704,7 @@ auto  GetFFTInputView([[maybe_unused]] OutputTensor &o,
     // FFT shorter than the size of the input signal. Create a new view of this
     // slice.
     if (act_fft_size < nom_fft_size) {
-      ends[RANK - 1] = nom_fft_size;
+      ends[RANK - 1] = act_fft_size;
       return i.Slice(starts, ends);
     }
     else { // FFT length is longer than the input. Pad input

diff --git a/test/00_transform/FFT.cu b/test/00_transform/FFT.cu
@@ -253,3 +253,39 @@ TYPED_TEST(FFTTestComplexTypes, IFFT2D16C2C)
   MATX_TEST_ASSERT_COMPARE(this->pb, avo, "a_out", this->thresh);
   MATX_EXIT_HANDLER();
 }
+
+
+TYPED_TEST(FFTTestComplexNonHalfTypes, FFT1D1024C2CShort)
+{
+  MATX_ENTER_HANDLER();
+  const index_t fft_dim = 1024;
+  this->pb->template InitAndRunTVGenerator<TypeParam>(
+      "00_transforms", "fft_operators", "fft_1d", {fft_dim, fft_dim - 16});
+
+  tensor_t<TypeParam, 1> av{{fft_dim}};
+  tensor_t<TypeParam, 1> avo{{fft_dim - 16}};
+  this->pb->NumpyToTensorView(av, "a_in");
+
+  fft(avo, av);
+  cudaStreamSynchronize(0);
+
+  MATX_TEST_ASSERT_COMPARE(this->pb, avo, "a_out", this->thresh);
+  MATX_EXIT_HANDLER();
+}
+
+TYPED_TEST(FFTTestComplexNonHalfTypes, IFFT1D1024C2CShort)
+{
+  MATX_ENTER_HANDLER();
+  const index_t fft_dim = 1024;
+  this->pb->template InitAndRunTVGenerator<TypeParam>(
+      "00_transforms", "fft_operators", "ifft_1d", {fft_dim, fft_dim - 16});
+  tensor_t<TypeParam, 1> av{{fft_dim}};
+  tensor_t<TypeParam, 1> avo{{fft_dim - 16}};
+  this->pb->NumpyToTensorView(av, "a_in");
+
+  ifft(avo, av);
+  cudaStreamSynchronize(0);
+
+  MATX_TEST_ASSERT_COMPARE(this->pb, avo, "a_out", this->thresh);
+  MATX_EXIT_HANDLER();
+}