Skip to content

Commit 9b2b9b5

Browse files
falseywinchnetfalseywinchnet
committed
Vectorize float32 spectrum transport
1 parent cdf0313 commit 9b2b9b5

1 file changed

Lines changed: 105 additions & 7 deletions

File tree

‎src/detail/bruun_kernel.hpp‎

Lines changed: 105 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1961,7 +1961,12 @@ class RFFT {
19611961
X[N / 2].im = 0.0f;
19621962

19631963
const int* RESTRICT kin = KINV.data();
1964-
for (int k = 1; k < N / 2; ++k) {
1964+
int k = 1;
1965+
for (; k + 3 < N / 2; k += 4) {
1966+
pack4_residues_to_complex_f32(X + k, work,
1967+
kin[k], kin[k + 1], kin[k + 2], kin[k + 3]);
1968+
}
1969+
for (; k < N / 2; ++k) {
19651970
const int m = kin[k];
19661971
X[k].re = work[2*m];
19671972
X[k].im = -work[2*m + 1];
@@ -1983,7 +1988,15 @@ class RFFT {
19831988
// Native heapopt order is a layout permutation; keep the output
19841989
// stream linear and pay the permutation on residue reads.
19851990
const int* RESTRICT native_leaf = NATIVE_LEAF.data();
1986-
for (int pos = 1; pos < N / 2; ++pos) {
1991+
int pos = 1;
1992+
for (; pos + 3 < N / 2; pos += 4) {
1993+
pack4_residues_to_complex_f32(X + pos, work,
1994+
native_leaf[pos],
1995+
native_leaf[pos + 1],
1996+
native_leaf[pos + 2],
1997+
native_leaf[pos + 3]);
1998+
}
1999+
for (; pos < N / 2; ++pos) {
19872000
const int m = native_leaf[pos];
19882001
X[pos].re = work[2*m];
19892002
X[pos].im = -work[2*m + 1];
@@ -2079,7 +2092,12 @@ class RFFT {
20792092
void inverse_f32(const complex_f32_t* RESTRICT X, float* RESTRICT out) const {
20802093
out[0] = 0.5f * (X[0].re + X[N / 2].re);
20812094
out[1] = 0.5f * (X[0].re - X[N / 2].re);
2082-
for (int m = 1; m < N / 2; ++m) {
2095+
int m = 1;
2096+
for (; m + 3 < N / 2; m += 4) {
2097+
unpack4_complex_to_residues_f32(out + 2*m, X,
2098+
IDX[m], IDX[m + 1], IDX[m + 2], IDX[m + 3]);
2099+
}
2100+
for (; m < N / 2; ++m) {
20832101
const int k = IDX[m];
20842102
out[2 * m] = X[k].re;
20852103
out[2 * m + 1] = -X[k].im;
@@ -2095,8 +2113,17 @@ class RFFT {
20952113
}
20962114
out[0] = 0.5f * (Xnative[0].re + Xnative[N / 2].re);
20972115
out[1] = 0.5f * (Xnative[0].re - Xnative[N / 2].re);
2098-
for (int pos = 1; pos < N / 2; ++pos) {
2099-
const int m = NATIVE_LEAF[pos];
2116+
const int* RESTRICT native_pos = NATIVE_POS.data();
2117+
int m = 1;
2118+
for (; m + 3 < N / 2; m += 4) {
2119+
unpack4_complex_to_residues_f32(out + 2*m, Xnative,
2120+
native_pos[m],
2121+
native_pos[m + 1],
2122+
native_pos[m + 2],
2123+
native_pos[m + 3]);
2124+
}
2125+
for (; m < N / 2; ++m) {
2126+
const int pos = native_pos[m];
21002127
out[2 * m] = Xnative[pos].re;
21012128
out[2 * m + 1] = -Xnative[pos].im;
21022129
}
@@ -2171,7 +2198,12 @@ class RFFT {
21712198
standardX[0] = nativeX[0];
21722199
standardX[N / 2] = nativeX[N / 2];
21732200
const int* RESTRICT map = STANDARD_NATIVE_POS.data();
2174-
for (int k = 1; k < N / 2; ++k) {
2201+
int k = 1;
2202+
for (; k + 3 < N / 2; k += 4) {
2203+
copy4_complex_f32(standardX + k, nativeX,
2204+
map[k], map[k + 1], map[k + 2], map[k + 3]);
2205+
}
2206+
for (; k < N / 2; ++k) {
21752207
standardX[k] = nativeX[map[k]];
21762208
}
21772209
#else
@@ -2189,7 +2221,12 @@ class RFFT {
21892221
nativeX[0] = standardX[0];
21902222
nativeX[N / 2] = standardX[N / 2];
21912223
const int* RESTRICT map = NATIVE_STANDARD_BIN.data();
2192-
for (int pos = 1; pos < N / 2; ++pos) {
2224+
int pos = 1;
2225+
for (; pos + 3 < N / 2; pos += 4) {
2226+
copy4_complex_f32(nativeX + pos, standardX,
2227+
map[pos], map[pos + 1], map[pos + 2], map[pos + 3]);
2228+
}
2229+
for (; pos < N / 2; ++pos) {
21932230
nativeX[pos] = standardX[map[pos]];
21942231
}
21952232
#else
@@ -2432,6 +2469,67 @@ class RFFT {
24322469
int m;
24332470
};
24342471

2472+
static inline void store4_interleaved_f32(float* RESTRICT dst,
2473+
float re0, float im0,
2474+
float re1, float im1,
2475+
float re2, float im2,
2476+
float re3, float im3) {
2477+
#if BRUUN_LEVEL >= 1
2478+
const bruun_v4f re = V4F_SET4(re0, re1, re2, re3);
2479+
const bruun_v4f im = V4F_SET4(im0, im1, im2, im3);
2480+
V4F_ST(dst, V4F_ZIPLO(re, im));
2481+
V4F_ST(dst + 4, V4F_ZIPHI(re, im));
2482+
#else
2483+
dst[0] = re0;
2484+
dst[1] = im0;
2485+
dst[2] = re1;
2486+
dst[3] = im1;
2487+
dst[4] = re2;
2488+
dst[5] = im2;
2489+
dst[6] = re3;
2490+
dst[7] = im3;
2491+
#endif
2492+
}
2493+
2494+
static inline void store4_complex_f32(complex_f32_t* RESTRICT dst,
2495+
float re0, float im0,
2496+
float re1, float im1,
2497+
float re2, float im2,
2498+
float re3, float im3) {
2499+
store4_interleaved_f32(reinterpret_cast<float*>(dst),
2500+
re0, im0, re1, im1, re2, im2, re3, im3);
2501+
}
2502+
2503+
static inline void pack4_residues_to_complex_f32(complex_f32_t* RESTRICT dst,
2504+
const float* RESTRICT residues,
2505+
int m0, int m1, int m2, int m3) {
2506+
store4_complex_f32(dst,
2507+
residues[2*m0], -residues[2*m0 + 1],
2508+
residues[2*m1], -residues[2*m1 + 1],
2509+
residues[2*m2], -residues[2*m2 + 1],
2510+
residues[2*m3], -residues[2*m3 + 1]);
2511+
}
2512+
2513+
static inline void copy4_complex_f32(complex_f32_t* RESTRICT dst,
2514+
const complex_f32_t* RESTRICT src,
2515+
int i0, int i1, int i2, int i3) {
2516+
store4_complex_f32(dst,
2517+
src[i0].re, src[i0].im,
2518+
src[i1].re, src[i1].im,
2519+
src[i2].re, src[i2].im,
2520+
src[i3].re, src[i3].im);
2521+
}
2522+
2523+
static inline void unpack4_complex_to_residues_f32(float* RESTRICT dst,
2524+
const complex_f32_t* RESTRICT src,
2525+
int i0, int i1, int i2, int i3) {
2526+
store4_interleaved_f32(dst,
2527+
src[i0].re, -src[i0].im,
2528+
src[i1].re, -src[i1].im,
2529+
src[i2].re, -src[i2].im,
2530+
src[i3].re, -src[i3].im);
2531+
}
2532+
24352533
static inline void norm_q1_fwd_f32(float* RESTRICT p, float c, float s) {
24362534
const float A0 = p[0];
24372535
const float B0 = p[1];

0 commit comments

Comments
 (0)