-
Notifications
You must be signed in to change notification settings - Fork 555
Expand file tree
/
Copy pathfftconvolve.cpp
More file actions
160 lines (132 loc) · 4.86 KB
/
Copy pathfftconvolve.cpp
File metadata and controls
160 lines (132 loc) · 4.86 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
/*******************************************************
* Copyright (c) 2022, ArrayFire
* All rights reserved.
*
* This file is distributed under 3-clause BSD license.
* The complete license agreement can be obtained at:
* http://arrayfire.com/licenses/BSD-3-Clause
********************************************************/
#include <fftconvolve.hpp>
#include <Array.hpp>
#include <common/dispatch.hpp>
#include <err_oneapi.hpp>
#include <fft.hpp>
#include <af/dim4.hpp>
#include <kernel/fftconvolve_common.hpp>
#include <kernel/fftconvolve_multiply.hpp>
#include <kernel/fftconvolve_pack.hpp>
#include <kernel/fftconvolve_pad.hpp>
#include <kernel/fftconvolve_reorder.hpp>
#include <cmath>
#include <type_traits>
#include <vector>
using af::dim4;
using std::ceil;
using std::conditional;
using std::is_integral;
using std::is_same;
using std::vector;
namespace arrayfire {
namespace oneapi {
template<typename T>
dim4 calcPackedSize(Array<T> const& i1, Array<T> const& i2, const dim_t rank) {
const dim4& i1d = i1.dims();
const dim4& i2d = i2.dims();
dim_t pd[4] = {1, 1, 1, 1};
// Pack both signal and filter on same memory array, this will ensure
// better use of batched cuFFT capabilities
pd[0] = nextpow2(static_cast<unsigned>(
static_cast<int>(ceil(i1d[0] / 2.f)) + i2d[0] - 1));
for (dim_t k = 1; k < rank; k++) {
pd[k] = nextpow2(static_cast<unsigned>(i1d[k] + i2d[k] - 1));
}
dim_t i1batch = 1;
dim_t i2batch = 1;
for (int k = rank; k < 4; k++) {
i1batch *= i1d[k];
i2batch *= i2d[k];
}
pd[rank] = (i1batch + i2batch);
return dim4(pd[0], pd[1], pd[2], pd[3]);
}
template<typename T>
Array<T> fftconvolve(Array<T> const& signal, Array<T> const& filter,
const bool expand, AF_BATCH_KIND kind, const int rank) {
using convT = typename conditional<is_integral<T>::value ||
is_same<T, float>::value ||
is_same<T, cfloat>::value,
float, double>::type;
using cT = typename conditional<is_same<convT, float>::value, cfloat,
cdouble>::type;
const dim4& sDims = signal.dims();
const dim4& fDims = filter.dims();
dim4 oDims(1);
if (expand) {
for (int d = 0; d < AF_MAX_DIMS; ++d) {
if (kind == AF_BATCH_NONE || kind == AF_BATCH_RHS) {
oDims[d] = sDims[d] + fDims[d] - 1;
} else {
oDims[d] = (d < rank ? sDims[d] + fDims[d] - 1 : sDims[d]);
}
}
} else {
oDims = sDims;
if (kind == AF_BATCH_RHS) {
for (int i = rank; i < AF_MAX_DIMS; ++i) { oDims[i] = fDims[i]; }
}
}
const dim4 pDims = calcPackedSize<T>(signal, filter, rank);
Array<cT> packed = createEmptyArray<cT>(pDims);
kernel::packDataHelper<cT, T>(packed, signal, filter, rank, kind);
kernel::padDataHelper<cT, T>(packed, signal, filter, rank, kind);
fft_inplace<cT>(packed, rank, true);
kernel::complexMultiplyHelper<cT, T>(packed, signal, filter, rank, kind);
// Compute inverse FFT only on complex-multiplied data
if (kind == AF_BATCH_RHS) {
vector<af_seq> seqs;
for (int k = 0; k < AF_MAX_DIMS; k++) {
if (k < rank) {
seqs.push_back({0., static_cast<double>(pDims[k] - 1), 1.});
} else if (k == rank) {
seqs.push_back({1., static_cast<double>(pDims[k] - 1), 1.});
} else {
seqs.push_back({0., 0., 1.});
}
}
Array<cT> subPacked = createSubArray<cT>(packed, seqs);
fft_inplace<cT>(subPacked, rank, false);
} else {
vector<af_seq> seqs;
for (int k = 0; k < AF_MAX_DIMS; k++) {
if (k < rank) {
seqs.push_back({0., static_cast<double>(pDims[k]) - 1, 1.});
} else if (k == rank) {
seqs.push_back({0., static_cast<double>(pDims[k] - 2), 1.});
} else {
seqs.push_back({0., 0., 1.});
}
}
Array<cT> subPacked = createSubArray<cT>(packed, seqs);
fft_inplace<cT>(subPacked, rank, false);
}
Array<T> out = createEmptyArray<T>(oDims);
kernel::reorderOutputHelper<T, cT>(out, packed, signal, filter, rank, kind,
expand);
return out;
}
#define INSTANTIATE(T) \
template Array<T> fftconvolve<T>(Array<T> const&, Array<T> const&, \
const bool, AF_BATCH_KIND, const int);
INSTANTIATE(double)
INSTANTIATE(float)
INSTANTIATE(uint)
INSTANTIATE(int)
INSTANTIATE(schar)
INSTANTIATE(uchar)
INSTANTIATE(char)
INSTANTIATE(uintl)
INSTANTIATE(intl)
INSTANTIATE(ushort)
INSTANTIATE(short)
} // namespace oneapi
} // namespace arrayfire