| 37 | |
| 38 | template<typename T> |
| 39 | Array<T> fftconvolve(Array<T> const& signal, Array<T> const& filter, |
| 40 | const bool expand, AF_BATCH_KIND kind, const int rank) { |
| 41 | using convT = typename std::conditional<std::is_integral<T>::value || |
| 42 | std::is_same<T, float>::value, |
| 43 | float, double>::type; |
| 44 | |
| 45 | constexpr bool IsTypeDouble = std::is_same<T, double>::value; |
| 46 | |
| 47 | const dim4& sd = signal.dims(); |
| 48 | const dim4& fd = filter.dims(); |
| 49 | dim_t fftScale = 1; |
| 50 | |
| 51 | dim4 packedDims(1, 1, 1, 1); |
| 52 | array<int, AF_MAX_DIMS> fftDims{}; // AF_MAX_DIMS(4) > rank |
| 53 | |
| 54 | // Pack both signal and filter on same memory array, this will ensure |
| 55 | // better use of batched FFT capabilities |
| 56 | fftDims[rank - 1] = nextpow2( |
| 57 | static_cast<unsigned>(static_cast<int>(ceil(sd[0] / 2.f)) + fd[0] - 1)); |
| 58 | packedDims[0] = 2 * fftDims[rank - 1]; |
| 59 | fftScale *= fftDims[rank - 1]; |
| 60 | |
| 61 | for (int k = 1; k < rank; k++) { |
| 62 | packedDims[k] = nextpow2(static_cast<unsigned>(sd[k] + fd[k] - 1)); |
| 63 | fftDims[rank - k - 1] = packedDims[k]; |
| 64 | fftScale *= fftDims[rank - k - 1]; |
| 65 | } |
| 66 | |
| 67 | dim_t sbatch = 1, fbatch = 1; |
| 68 | for (int k = rank; k < AF_MAX_DIMS; k++) { |
| 69 | sbatch *= sd[k]; |
| 70 | fbatch *= fd[k]; |
| 71 | } |
| 72 | packedDims[rank] = (sbatch + fbatch); |
| 73 | |
| 74 | Array<convT> packed = createEmptyArray<convT>(packedDims); |
| 75 | |
| 76 | dim4 paddedSigDims(packedDims[0], (1 < rank ? packedDims[1] : sd[1]), |
| 77 | (2 < rank ? packedDims[2] : sd[2]), |
| 78 | (3 < rank ? packedDims[3] : sd[3])); |
| 79 | dim4 paddedFilDims(packedDims[0], (1 < rank ? packedDims[1] : fd[1]), |
| 80 | (2 < rank ? packedDims[2] : fd[2]), |
| 81 | (3 < rank ? packedDims[3] : fd[3])); |
| 82 | dim4 paddedSigStrides = calcStrides(paddedSigDims); |
| 83 | dim4 paddedFilStrides = calcStrides(paddedFilDims); |
| 84 | |
| 85 | // Number of packed complex elements in dimension 0 |
| 86 | dim_t sig_half_d0 = divup(sd[0], 2); |
| 87 | |
| 88 | // Pack signal in a complex matrix where first dimension is half the input |
| 89 | // (allows faster FFT computation) and pad array to a power of 2 with 0s |
| 90 | getQueue().enqueue(kernel::packData<convT, T>, packed, paddedSigDims, |
| 91 | paddedSigStrides, signal); |
| 92 | |
| 93 | // Pad filter array with 0s |
| 94 | const dim_t offset = paddedSigStrides[3] * paddedSigDims[3]; |
| 95 | getQueue().enqueue(kernel::padArray<convT, T>, packed, paddedFilDims, |
| 96 | paddedFilStrides, filter, offset); |