heffte_stock_complex.h

/*
    -- heFFTe --
       Univ. of Tennessee, Knoxville
       @date
*/
 
#ifndef HEFFTE_STOCK_COMPLEX_H
#define HEFFTE_STOCK_COMPLEX_H
 
#include <type_traits>
#include <iostream>
 
#include "heffte_stock_vec_types.h"
 
namespace heffte {
namespace stock {
template<typename F, int L>
class alignas(L*sizeof(F)) Complex {
    public:
        // One 64-bit Complex-- 2 doubles-- pack<double, 2>::type == _m128d
        // Two 64-bit Complex-- 4 doubles-- pack<double, 4>::type == _m256d
        // Two 64-bit Complex-- 4 floats -- pack<float, 4>::type == _m128
        // Four 64-bit Complex-- 8 floats -- pack<float, 8>::type == _m256
 
        // Constructors for the Complex class
        explicit Complex(F* const f): var(mm_load<F,L>(f)) {}
        explicit Complex(std::initializer_list<F> il): var(mm_load<F,L>(il.begin())) {};
        explicit Complex(typename pack<F,L>::type v): var(v) {}
        explicit Complex(F x, F y): var(mm_pair_set<F,L>(x, y)) {}
        explicit Complex(std::complex<F> const *c): var(mm_complex_load<F,L>(c)) {}
        explicit Complex(std::complex<F> const *c, int stride): var(mm_complex_load<F,L>(c, stride)) {}
        explicit Complex(std::initializer_list<std::complex<F>> il): var(il.size() == 1 ? mm_pair_set<F,L>((*il.begin()).real(), (*il.begin()).imag()) : mm_complex_load<F,L>(il.begin())) {};
        explicit Complex(): var(mm_zero<F,L>()) {}
 
        /* Basic operations with another pack of complex numbers */
 
        Complex<F,L> operator+(Complex<F,L> const &o) {
            return Complex(mm_add(var, o.var));
        }
 
        Complex<F,L> operator-(Complex<F,L> const &o) {
            return Complex(mm_sub(var, o.var));
        }
 
        Complex<F,L> operator*(Complex<F,L> const & o) {
            return Complex(mm_complex_mul(var, o.var));
        }
 
        Complex<F,L> operator/(Complex<F,L> const & o) {
            return Complex(mm_complex_div(var, o.var));
        }
 
        Complex<F,L> operator+=(Complex<F,L> const &o) {
            var = mm_add(var, o.var);
            return *this;
        }
 
        Complex<F,L> operator-=(Complex<F,L> const &o) {
            var = mm_sub(var, o.var);
            return *this;
        }
 
        Complex<F,L> operator*=(Complex<F,L> const &o) {
            var = mm_complex_mul(var, o.var);
            return *this;
        }
 
        Complex<F,L> operator/=(Complex<F,L> const &o) {
            var = mm_complex_div(var, o.var);
            return *this;
        }
 
        /* Basic operations with a single real floating point number */
 
        Complex<F,L> operator+(F o) {
            return Complex(mm_add(var, mm_pair_set<F,L>(o, 0)));
        }
 
        Complex<F,L> operator-(F o) {
            return Complex(mm_sub(var, mm_pair_set<F,L>(o, 0)));
        }
 
        Complex<F,L> operator*(F o) {
            return Complex(mm_mul(var, mm_set1<F,L>(o)));
        }
 
        Complex<F,L> operator/(F o) {
            return Complex(mm_div(var, mm_set1<F,L>(o)));
        }
 
        Complex<F,L> operator+=(F o) {
            var = mm_add(var, mm_pair_set<F,L>(o, 0));
            return *this;
        }
 
        Complex<F,L> operator-=(F o) {
            var = mm_sub(var, mm_pair_set<F,L>(o, 0));
            return *this;
        }
 
        Complex<F,L> operator*=(F o) {
            var = mm_mul(var, mm_set1<F,L>(o));
            return *this;
        }
 
        // Divide by and assign a floating point number
        Complex<F,L> operator/=(F o) {
            var = mm_div(var, mm_set1<F,L>(o));
            return *this;
        }
 
        /* Other methods */
 
        Complex<F,L> fmadd(Complex<F,L> const & y, Complex<F,L> const & z) {
            return Complex(mm_complex_fmadd(var, y.var, z.var));
        }
 
        Complex<F,L> fmsub(Complex<F,L> const & y, Complex<F,L> const & z) {
            return Complex(mm_complex_fmsub(var, y.var, z.var));
        }
 
        Complex<F,L> operator-() {
            return Complex<F,L>(mm_neg(var));
        }
 
        void modulus(F* dest) {
            typename pack<F, L>::type res = mm_complex_mod(var);
            for(int i = 0; i < L/2; i++) {
                dest[i] = res[i*2];
            }
        }
 
        typename pack<F,L>::type modulus() {
            return mm_complex_mod(var);
        }
 
        Complex<F,L> conjugate() {
            return Complex(mm_complex_conj(var));
        }
 
        Complex<F,L> __mul_i() {
            return Complex(mm_complex_mul_i(var));
        }
 
        Complex<F,L> __mul_neg_i() {
            return Complex(mm_complex_mul_neg_i(var));
        }
 
        void get(F* dest) {
            mm_store<F,L>(dest, var);
        }
 
        std::complex<F> operator[](std::size_t idx) {
            return std::complex<F>(reinterpret_cast<F*>(&var)[2*idx], reinterpret_cast<F*>(&var)[2*idx + 1]);
        }
 
        typename pack<F,L>::type get() const {
            return var;
        }
 
    private:
        typename pack<F,L>::type var;
};
 
 
 
 
template<typename F, int L>
inline std::ostream& operator<<(std::ostream& os, const Complex<F,L>& dt){
    typename pack<F, L>::type var = dt.get();
    os << "( ";
    for(int i = 0; i < L; i+=2) {
        os << var[i];
        if(var[i+1] < 0) os << " - " << -var[i+1] << "i";
        else os << " + " << var[i+1] << "i";
        if(i+2 < L) os << ", ";
    }
    os << " )";
    return os;
}
 
}
}
 
#endif // HEFFTE_STOCK_COMPLEX_H