local_multiply Class — pytorch Architecture

Source Code

aten/src/ATen/test/vec_test_all_types.h lines 1283–1312

template <typename T>
std::enable_if_t<is_complex<Complex<T>>::value, Complex<T>> local_multiply(Complex<T> x, Complex<T> y) {
#if defined(TEST_AGAINST_DEFAULT)
    return x * y;
#else
    //(a + bi)  * (c + di) = (ac - bd) + (ad + bc)i
    T x_real = x.real();
    T x_imag = x.imag();
    T y_real = y.real();
    T y_imag = y.imag();
#if defined(CPU_CAPABILITY_VSX) || defined(CPU_CAPABILITY_ZVECTOR)
    //check multiplication considering swap and fma
    T rr = x_real * y_real;
    T ii = x_imag * y_real;
    T neg_imag = -y_imag;
    rr = fma(x_imag, neg_imag, rr);
    ii = fma(x_real, y_imag, ii);
#else
    // replicate order
    PreventFma noFma;
    T ac = x_real * y_real;
    T bd = x_imag * y_imag;
    T ad = x_real * y_imag;
    T bc = x_imag * (-y_real);
    T rr = noFma.sub(ac, bd);
    T ii = noFma.sub(ad, bc);
#endif
    return Complex<T>(rr, ii);
#endif
}