path: root/library/cpp/dot_product/dot_product_sse.cpp

#include "dot_product_sse.h"

#include <library/cpp/sse/sse.h>
#include <util/system/platform.h>
#include <util/system/compiler.h>

#ifdef ARCADIA_SSE
i32 DotProductSse(const i8* lhs, const i8* rhs, size_t length) noexcept {
    const __m128i zero = _mm_setzero_si128();
    __m128i resVec = zero;
    while (length >= 16) {
        __m128i lVec = _mm_loadu_si128((const __m128i*)lhs);
        __m128i rVec = _mm_loadu_si128((const __m128i*)rhs);

#ifdef _sse4_1_
        __m128i lLo = _mm_cvtepi8_epi16(lVec);
        __m128i rLo = _mm_cvtepi8_epi16(rVec);
        __m128i lHi = _mm_cvtepi8_epi16(_mm_alignr_epi8(lVec, lVec, 8));
        __m128i rHi = _mm_cvtepi8_epi16(_mm_alignr_epi8(rVec, rVec, 8));
#else
        __m128i lLo = _mm_srai_epi16(_mm_unpacklo_epi8(zero, lVec), 8);
        __m128i rLo = _mm_srai_epi16(_mm_unpacklo_epi8(zero, rVec), 8);
        __m128i lHi = _mm_srai_epi16(_mm_unpackhi_epi8(zero, lVec), 8);
        __m128i rHi = _mm_srai_epi16(_mm_unpackhi_epi8(zero, rVec), 8);
#endif
        resVec = _mm_add_epi32(resVec,
                               _mm_add_epi32(_mm_madd_epi16(lLo, rLo), _mm_madd_epi16(lHi, rHi)));

        lhs += 16;
        rhs += 16;
        length -= 16;
    }

    alignas(16) i32 res[4];
    _mm_store_si128((__m128i*)res, resVec);
    i32 sum = res[0] + res[1] + res[2] + res[3];
    for (size_t i = 0; i < length; ++i) {
        sum += static_cast<i32>(lhs[i]) * static_cast<i32>(rhs[i]);
    }

    return sum;
}

ui32 DotProductSse(const ui8* lhs, const ui8* rhs, size_t length) noexcept {
    const __m128i zero = _mm_setzero_si128();
    __m128i resVec = zero;
    while (length >= 16) {
        __m128i lVec = _mm_loadu_si128((const __m128i*)lhs);
        __m128i rVec = _mm_loadu_si128((const __m128i*)rhs);

        __m128i lLo = _mm_unpacklo_epi8(lVec, zero);
        __m128i rLo = _mm_unpacklo_epi8(rVec, zero);
        __m128i lHi = _mm_unpackhi_epi8(lVec, zero);
        __m128i rHi = _mm_unpackhi_epi8(rVec, zero);

        resVec = _mm_add_epi32(resVec,
                               _mm_add_epi32(_mm_madd_epi16(lLo, rLo), _mm_madd_epi16(lHi, rHi)));

        lhs += 16;
        rhs += 16;
        length -= 16;
    }

    alignas(16) i32 res[4];
    _mm_store_si128((__m128i*)res, resVec);
    i32 sum = res[0] + res[1] + res[2] + res[3];
    for (size_t i = 0; i < length; ++i) {
        sum += static_cast<i32>(lhs[i]) * static_cast<i32>(rhs[i]);
    }

    return static_cast<ui32>(sum);
}
#ifdef _sse4_1_

i64 DotProductSse(const i32* lhs, const i32* rhs, size_t length) noexcept {
    __m128i zero = _mm_setzero_si128();
    __m128i res = zero;

    while (length >= 4) {
        __m128i a = _mm_loadu_si128((const __m128i*)lhs);
        __m128i b = _mm_loadu_si128((const __m128i*)rhs);
        res = _mm_add_epi64(_mm_mul_epi32(a, b), res);    // This is lower parts multiplication
        a = _mm_alignr_epi8(a, a, 4);
        b = _mm_alignr_epi8(b, b, 4);
        res = _mm_add_epi64(_mm_mul_epi32(a, b), res);
        rhs += 4;
        lhs += 4;
        length -= 4;
    }

    alignas(16) i64 r[2];
    _mm_store_si128((__m128i*)r, res);
    i64 sum = r[0] + r[1];

    for (size_t i = 0; i < length; ++i) {
        sum += static_cast<i64>(lhs[i]) * static_cast<i64>(rhs[i]);
    }

    return sum;
}

#else
#include "dot_product_simple.h"

i64 DotProductSse(const i32* lhs, const i32* rhs, size_t length) noexcept {
    return DotProductSimple(lhs, rhs, length);
}

#endif

float DotProductSse(const float* lhs, const float* rhs, size_t length) noexcept {
    __m128 sum1 = _mm_setzero_ps();
    __m128 sum2 = _mm_setzero_ps();
    __m128 a1, b1, a2, b2, m1, m2;

    while (length >= 8) {
        a1 = _mm_loadu_ps(lhs);
        b1 = _mm_loadu_ps(rhs);
        m1 = _mm_mul_ps(a1, b1);

        a2 = _mm_loadu_ps(lhs + 4);
        sum1 = _mm_add_ps(sum1, m1);

        b2 = _mm_loadu_ps(rhs + 4);
        m2 = _mm_mul_ps(a2, b2);

        sum2 = _mm_add_ps(sum2, m2);

        length -= 8;
        lhs += 8;
        rhs += 8;
    }

    if (length >= 4) {
        a1 = _mm_loadu_ps(lhs);
        b1 = _mm_loadu_ps(rhs);
        sum1 = _mm_add_ps(sum1, _mm_mul_ps(a1, b1));

        length -= 4;
        lhs += 4;
        rhs += 4;
    }

    sum1 = _mm_add_ps(sum1, sum2);

    if (length) {
        switch (length) {
            case 3:
                a1 = _mm_set_ps(0.0f, lhs[2], lhs[1], lhs[0]);
                b1 = _mm_set_ps(0.0f, rhs[2], rhs[1], rhs[0]);
                break;

            case 2:
                a1 = _mm_set_ps(0.0f, 0.0f, lhs[1], lhs[0]);
                b1 = _mm_set_ps(0.0f, 0.0f, rhs[1], rhs[0]);
                break;

            case 1:
                a1 = _mm_set_ps(0.0f, 0.0f, 0.0f, lhs[0]);
                b1 = _mm_set_ps(0.0f, 0.0f, 0.0f, rhs[0]);
                break;

            default:
                Y_UNREACHABLE();
        }

        sum1 = _mm_add_ps(sum1, _mm_mul_ps(a1, b1));
    }

    alignas(16) float res[4];
    _mm_store_ps(res, sum1);

    return res[0] + res[1] + res[2] + res[3];
}

double DotProductSse(const double* lhs, const double* rhs, size_t length) noexcept {
    __m128d sum1 = _mm_setzero_pd();
    __m128d sum2 = _mm_setzero_pd();
    __m128d a1, b1, a2, b2;

    while (length >= 4) {
        a1 = _mm_loadu_pd(lhs);
        b1 = _mm_loadu_pd(rhs);
        sum1 = _mm_add_pd(sum1, _mm_mul_pd(a1, b1));

        a2 = _mm_loadu_pd(lhs + 2);
        b2 = _mm_loadu_pd(rhs + 2);
        sum2 = _mm_add_pd(sum2, _mm_mul_pd(a2, b2));

        length -= 4;
        lhs += 4;
        rhs += 4;
    }

    if (length >= 2) {
        a1 = _mm_loadu_pd(lhs);
        b1 = _mm_loadu_pd(rhs);
        sum1 = _mm_add_pd(sum1, _mm_mul_pd(a1, b1));

        length -= 2;
        lhs += 2;
        rhs += 2;
    }

    sum1 = _mm_add_pd(sum1, sum2);

    if (length > 0) {
        a1 = _mm_set_pd(lhs[0], 0.0);
        b1 = _mm_set_pd(rhs[0], 0.0);
        sum1 = _mm_add_pd(sum1, _mm_mul_pd(a1, b1));
    }

    alignas(16) double res[2];
    _mm_store_pd(res, sum1);

    return res[0] + res[1];
}

template <bool computeRR>
Y_FORCE_INLINE
static void TriWayDotProductIterationSse(__m128& sumLL, __m128& sumLR, __m128& sumRR, const __m128 a, const __m128 b) {
    sumLL = _mm_add_ps(sumLL, _mm_mul_ps(a, a));
    sumLR = _mm_add_ps(sumLR, _mm_mul_ps(a, b));

    if constexpr (computeRR) {
        sumRR = _mm_add_ps(sumRR, _mm_mul_ps(b, b));
    }
}

template <bool computeRR>
TTriWayDotProduct<float> TriWayDotProductSseImpl(
    const float* lhs,
    const float* rhs,
    size_t length) noexcept
{
    __m128 sumLL1 = _mm_setzero_ps();
    __m128 sumLR1 = _mm_setzero_ps();
    __m128 sumRR1 = _mm_setzero_ps();
    __m128 sumLL2 = _mm_setzero_ps();
    __m128 sumLR2 = _mm_setzero_ps();
    __m128 sumRR2 = _mm_setzero_ps();

    while (length >= 8) {
        TriWayDotProductIterationSse<computeRR>(sumLL1, sumLR1, sumRR1, _mm_loadu_ps(lhs + 0), _mm_loadu_ps(rhs + 0));
        TriWayDotProductIterationSse<computeRR>(sumLL2, sumLR2, sumRR2, _mm_loadu_ps(lhs + 4), _mm_loadu_ps(rhs + 4));
        length -= 8;
        lhs += 8;
        rhs += 8;
    }

    if (length >= 4) {
        TriWayDotProductIterationSse<computeRR>(sumLL1, sumLR1, sumRR1, _mm_loadu_ps(lhs + 0), _mm_loadu_ps(rhs + 0));
        length -= 4;
        lhs += 4;
        rhs += 4;
    }

    sumLL1 = _mm_add_ps(sumLL1, sumLL2);
    sumLR1 = _mm_add_ps(sumLR1, sumLR2);

    if (computeRR) {
        sumRR1 = _mm_add_ps(sumRR1, sumRR2);
    }

    if (length) {
        __m128 a, b;
        switch (length) {
            case 3:
                a = _mm_set_ps(0.0f, lhs[2], lhs[1], lhs[0]);
                b = _mm_set_ps(0.0f, rhs[2], rhs[1], rhs[0]);
                break;
            case 2:
                a = _mm_set_ps(0.0f, 0.0f, lhs[1], lhs[0]);
                b = _mm_set_ps(0.0f, 0.0f, rhs[1], rhs[0]);
                break;
            case 1:
                a = _mm_set_ps(0.0f, 0.0f, 0.0f, lhs[0]);
                b = _mm_set_ps(0.0f, 0.0f, 0.0f, rhs[0]);
                break;
            default:
                Y_UNREACHABLE();
        }
        TriWayDotProductIterationSse<computeRR>(sumLL1, sumLR1, sumRR1, a, b);
    }

    __m128 t0 = sumLL1;
    __m128 t1 = sumLR1;
    __m128 t2 = sumRR1;
    __m128 t3 = _mm_setzero_ps();
    _MM_TRANSPOSE4_PS(t0, t1, t2, t3);
    t0 = _mm_add_ps(t0, t1);
    t0 = _mm_add_ps(t0, t2);
    t0 = _mm_add_ps(t0, t3);

    alignas(16) float res[4];
    _mm_store_ps(res, t0);
    TTriWayDotProduct<float> result{res[0], res[1], res[2]};
    if (!computeRR) {
        static constexpr const TTriWayDotProduct<float> def;
        result.RR = def.RR;
    }
    return result;
}

TTriWayDotProduct<float> TriWayDotProductSse(
    const float* lhs,
    const float* rhs,
    size_t length,
    bool computeRR) noexcept
{
    if (computeRR) {
        return TriWayDotProductSseImpl<true>(lhs, rhs, length);
    } else {
        return TriWayDotProductSseImpl<false>(lhs, rhs, length);
    }
}

#endif // ARCADIA_SSE