203 lines
6.2 KiB
C++
203 lines
6.2 KiB
C++
#include <algorithm>
|
|
#include <cmath>
|
|
#include <vector>
|
|
|
|
#include <immintrin.h>
|
|
|
|
#include "trigdx/lookup_avx.hpp"
|
|
|
|
template <std::size_t NR_SAMPLES> struct LookupAVXBackend<NR_SAMPLES>::Impl {
|
|
std::vector<float> lookup;
|
|
static constexpr std::size_t MASK = NR_SAMPLES - 1;
|
|
static constexpr float SCALE = NR_SAMPLES / (2.0f * float(M_PI));
|
|
|
|
void init() {
|
|
lookup.resize(NR_SAMPLES);
|
|
for (std::size_t i = 0; i < NR_SAMPLES; ++i) {
|
|
lookup[i] = sinf(i * (2.0f * float(M_PI) / NR_SAMPLES));
|
|
}
|
|
}
|
|
|
|
void compute_sincosf(std::size_t n, const float *x, float *s,
|
|
float *c) const {
|
|
#if defined(__AVX__)
|
|
constexpr std::size_t VL = 8; // AVX processes 8 floats
|
|
const __m256 scale = _mm256_set1_ps(SCALE);
|
|
const __m256i mask = _mm256_set1_epi32(MASK);
|
|
const __m256i quarter_pi = _mm256_set1_epi32(NR_SAMPLES / 4);
|
|
|
|
std::size_t i = 0;
|
|
for (; i + VL <= n; i += VL) {
|
|
__m256 vx = _mm256_loadu_ps(&x[i]);
|
|
__m256 scaled = _mm256_mul_ps(vx, scale);
|
|
__m256i idx = _mm256_cvtps_epi32(scaled);
|
|
__m256i idx_cos = _mm256_add_epi32(idx, quarter_pi);
|
|
|
|
idx = _mm256_and_si256(idx, mask);
|
|
idx_cos = _mm256_and_si256(idx_cos, mask);
|
|
|
|
#if defined(__AVX2__)
|
|
__m256 sinv = _mm256_i32gather_ps(lookup.data(), idx, 4);
|
|
__m256 cosv = _mm256_i32gather_ps(lookup.data(), idx_cos, 4);
|
|
#else
|
|
// fallback gather for AVX1
|
|
float sin_tmp[VL], cos_tmp[VL];
|
|
int idx_a[VL], idxc_a[VL];
|
|
_mm256_store_si256((__m256i *)idx_a, idx);
|
|
_mm256_store_si256((__m256i *)idxc_a, idx_cos);
|
|
for (std::size_t k = 0; k < VL; ++k) {
|
|
sin_tmp[k] = lookup[idx_a[k]];
|
|
cos_tmp[k] = lookup[idxc_a[k]];
|
|
}
|
|
__m256 sinv = _mm256_load_ps(sin_tmp);
|
|
__m256 cosv = _mm256_load_ps(cos_tmp);
|
|
#endif
|
|
_mm256_storeu_ps(&s[i], sinv);
|
|
_mm256_storeu_ps(&c[i], cosv);
|
|
}
|
|
|
|
// scalar remainder
|
|
for (; i < n; ++i) {
|
|
std::size_t idx = static_cast<std::size_t>(x[i] * SCALE) & MASK;
|
|
std::size_t idx_cos = (idx + NR_SAMPLES / 4) & MASK;
|
|
s[i] = lookup[idx];
|
|
c[i] = lookup[idx_cos];
|
|
}
|
|
#else
|
|
// No AVX: scalar path
|
|
for (std::size_t i = 0; i < n; ++i) {
|
|
std::size_t idx = static_cast<std::size_t>(x[i] * SCALE) & MASK;
|
|
std::size_t idx_cos = (idx + NR_SAMPLES / 4) & MASK;
|
|
s[i] = lookup[idx];
|
|
c[i] = lookup[idx_cos];
|
|
}
|
|
#endif
|
|
}
|
|
|
|
void compute_sinf(std::size_t n, const float *x, float *s) const {
|
|
#if defined(__AVX__)
|
|
constexpr std::size_t VL = 8; // AVX processes 8 floats
|
|
const __m256 scale = _mm256_set1_ps(SCALE);
|
|
const __m256i mask = _mm256_set1_epi32(MASK);
|
|
const __m256i quarter_pi = _mm256_set1_epi32(NR_SAMPLES / 4);
|
|
|
|
std::size_t i = 0;
|
|
for (; i + VL <= n; i += VL) {
|
|
__m256 vx = _mm256_loadu_ps(&x[i]);
|
|
__m256 scaled = _mm256_mul_ps(vx, scale);
|
|
__m256i idx = _mm256_cvtps_epi32(scaled);
|
|
|
|
idx = _mm256_and_si256(idx, mask);
|
|
|
|
#if defined(__AVX2__)
|
|
__m256 sinv = _mm256_i32gather_ps(lookup.data(), idx, 4);
|
|
#else
|
|
// fallback gather for AVX1
|
|
float sin_tmp[VL];
|
|
int idx_a[VL], idxc_a[VL];
|
|
_mm256_store_si256((__m256i *)idx_a, idx);
|
|
for (std::size_t k = 0; k < VL; ++k) {
|
|
sin_tmp[k] = lookup[idx_a[k]];
|
|
}
|
|
__m256 sinv = _mm256_load_ps(sin_tmp);
|
|
#endif
|
|
_mm256_storeu_ps(&s[i], sinv);
|
|
}
|
|
|
|
// scalar remainder
|
|
for (; i < n; ++i) {
|
|
std::size_t idx = static_cast<std::size_t>(x[i] * SCALE) & MASK;
|
|
s[i] = lookup[idx];
|
|
}
|
|
#else
|
|
// No AVX: scalar path
|
|
for (std::size_t i = 0; i < n; ++i) {
|
|
std::size_t idx = static_cast<std::size_t>(x[i] * SCALE) & MASK;
|
|
s[i] = lookup[idx];
|
|
}
|
|
#endif
|
|
}
|
|
|
|
void compute_cosf(std::size_t n, const float *x, float *c) const {
|
|
#if defined(__AVX__)
|
|
constexpr std::size_t VL = 8; // AVX processes 8 floats
|
|
const __m256 scale = _mm256_set1_ps(SCALE);
|
|
const __m256i mask = _mm256_set1_epi32(MASK);
|
|
const __m256i quarter_pi = _mm256_set1_epi32(NR_SAMPLES / 4);
|
|
|
|
std::size_t i = 0;
|
|
for (; i + VL <= n; i += VL) {
|
|
__m256 vx = _mm256_loadu_ps(&x[i]);
|
|
__m256 scaled = _mm256_mul_ps(vx, scale);
|
|
__m256i idx = _mm256_cvtps_epi32(scaled);
|
|
__m256i idx_cos = _mm256_add_epi32(idx, quarter_pi);
|
|
|
|
idx_cos = _mm256_and_si256(idx_cos, mask);
|
|
|
|
#if defined(__AVX2__)
|
|
__m256 cosv = _mm256_i32gather_ps(lookup.data(), idx_cos, 4);
|
|
#else
|
|
// fallback gather for AVX1
|
|
float cos_tmp[VL];
|
|
int idxc_a[VL];
|
|
_mm256_store_si256((__m256i *)idxc_a, idx_cos);
|
|
for (std::size_t k = 0; k < VL; ++k) {
|
|
cos_tmp[k] = lookup[idxc_a[k]];
|
|
}
|
|
__m256 cosv = _mm256_load_ps(cos_tmp);
|
|
#endif
|
|
_mm256_storeu_ps(&c[i], cosv);
|
|
}
|
|
|
|
// scalar remainder
|
|
for (; i < n; ++i) {
|
|
std::size_t idx = static_cast<std::size_t>(x[i] * SCALE) & MASK;
|
|
std::size_t idx_cos = (idx + NR_SAMPLES / 4) & MASK;
|
|
c[i] = lookup[idx_cos];
|
|
}
|
|
#else
|
|
// No AVX: scalar path
|
|
for (std::size_t i = 0; i < n; ++i) {
|
|
std::size_t idx = static_cast<std::size_t>(x[i] * SCALE) & MASK;
|
|
std::size_t idx_cos = (idx + NR_SAMPLES / 4) & MASK;
|
|
c[i] = lookup[idx_cos];
|
|
}
|
|
#endif
|
|
}
|
|
};
|
|
|
|
template <std::size_t NR_SAMPLES>
|
|
LookupAVXBackend<NR_SAMPLES>::LookupAVXBackend()
|
|
: impl(std::make_unique<Impl>()) {}
|
|
|
|
template <std::size_t NR_SAMPLES>
|
|
LookupAVXBackend<NR_SAMPLES>::~LookupAVXBackend() = default;
|
|
|
|
template <std::size_t NR_SAMPLES>
|
|
void LookupAVXBackend<NR_SAMPLES>::init(size_t) {
|
|
impl->init();
|
|
}
|
|
|
|
template <std::size_t NR_SAMPLES>
|
|
void LookupAVXBackend<NR_SAMPLES>::compute_sinf(std::size_t n, const float *x,
|
|
float *s) const {
|
|
impl->compute_sinf(n, x, s);
|
|
}
|
|
|
|
template <std::size_t NR_SAMPLES>
|
|
void LookupAVXBackend<NR_SAMPLES>::compute_cosf(std::size_t n, const float *x,
|
|
float *c) const {
|
|
impl->compute_cosf(n, x, c);
|
|
}
|
|
|
|
template <std::size_t NR_SAMPLES>
|
|
void LookupAVXBackend<NR_SAMPLES>::compute_sincosf(std::size_t n,
|
|
const float *x, float *s,
|
|
float *c) const {
|
|
impl->compute_sincosf(n, x, s, c);
|
|
}
|
|
|
|
// Explicit instantiations
|
|
template class LookupAVXBackend<16384>;
|
|
template class LookupAVXBackend<32768>;
|