// Copyright (C) 2018-2020 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "cpu_x86_sse42/blob_transform_sse42.hpp"

#include <nmmintrin.h>  // SSE 4.2

namespace InferenceEngine {

//------------------------------------------------------------------------
//
// OpenCV universal intrinsics (refactored), x86 specific - require SSE4.1
//
//------------------------------------------------------------------------

static inline void mm_load_deinterleave(const uint8_t* ptr, __m128i& a, __m128i& b, __m128i& c) {
    const __m128i m0 = _mm_setr_epi8(0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0);
    const __m128i m1 = _mm_setr_epi8(0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0);
    __m128i s0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(ptr));
    __m128i s1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(ptr + 16));
    __m128i s2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(ptr + 32));
    __m128i a0 = _mm_blendv_epi8(_mm_blendv_epi8(s0, s1, m0), s2, m1);
    __m128i b0 = _mm_blendv_epi8(_mm_blendv_epi8(s1, s2, m0), s0, m1);
    __m128i c0 = _mm_blendv_epi8(_mm_blendv_epi8(s2, s0, m0), s1, m1);
    const __m128i sh_b = _mm_setr_epi8(0, 3, 6, 9, 12, 15, 2, 5, 8, 11, 14, 1, 4, 7, 10, 13);
    const __m128i sh_g = _mm_setr_epi8(1, 4, 7, 10, 13, 0, 3, 6, 9, 12, 15, 2, 5, 8, 11, 14);
    const __m128i sh_r = _mm_setr_epi8(2, 5, 8, 11, 14, 1, 4, 7, 10, 13, 0, 3, 6, 9, 12, 15);
    a0 = _mm_shuffle_epi8(a0, sh_b);
    b0 = _mm_shuffle_epi8(b0, sh_g);
    c0 = _mm_shuffle_epi8(c0, sh_r);
    a = a0;
    b = b0;
    c = c0;
}

static inline void mm_load_deinterleave(const float* ptr, __m128& a, __m128& b, __m128& c) {
    __m128 t0 = _mm_loadu_ps(ptr + 0);
    __m128 t1 = _mm_loadu_ps(ptr + 4);
    __m128 t2 = _mm_loadu_ps(ptr + 8);

    __m128 at12 = _mm_shuffle_ps(t1, t2, _MM_SHUFFLE(0, 1, 0, 2));
    a = _mm_shuffle_ps(t0, at12, _MM_SHUFFLE(2, 0, 3, 0));

    __m128 bt01 = _mm_shuffle_ps(t0, t1, _MM_SHUFFLE(0, 0, 0, 1));
    __m128 bt12 = _mm_shuffle_ps(t1, t2, _MM_SHUFFLE(0, 2, 0, 3));
    b = _mm_shuffle_ps(bt01, bt12, _MM_SHUFFLE(2, 0, 2, 0));

    __m128 ct01 = _mm_shuffle_ps(t0, t1, _MM_SHUFFLE(0, 1, 0, 2));
    c = _mm_shuffle_ps(ct01, t2, _MM_SHUFFLE(3, 0, 2, 0));
}

static inline void mm_store_interleave(uint8_t* ptr, __m128i a, __m128i b, __m128i c) {
    const __m128i sh_a = _mm_setr_epi8(0, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14, 9, 4, 15, 10, 5);
    const __m128i sh_b = _mm_setr_epi8(5, 0, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14, 9, 4, 15, 10);
    const __m128i sh_c = _mm_setr_epi8(10, 5, 0, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14, 9, 4, 15);
    __m128i a0 = _mm_shuffle_epi8(a, sh_a);
    __m128i b0 = _mm_shuffle_epi8(b, sh_b);
    __m128i c0 = _mm_shuffle_epi8(c, sh_c);

    const __m128i m0 = _mm_setr_epi8(0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0);
    const __m128i m1 = _mm_setr_epi8(0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0);
    __m128i v0 = _mm_blendv_epi8(_mm_blendv_epi8(a0, b0, m1), c0, m0);
    __m128i v1 = _mm_blendv_epi8(_mm_blendv_epi8(b0, c0, m1), a0, m0);
    __m128i v2 = _mm_blendv_epi8(_mm_blendv_epi8(c0, a0, m1), b0, m0);

    _mm_storeu_si128(reinterpret_cast<__m128i*>(ptr), v0);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(ptr + 16), v1);
    _mm_storeu_si128(reinterpret_cast<__m128i*>(ptr + 32), v2);
}

static inline void mm_store_interleave(float* ptr, __m128 a, __m128 b, __m128 c) {
    __m128 u0 = _mm_shuffle_ps(a, b, _MM_SHUFFLE(0, 0, 0, 0));
    __m128 u1 = _mm_shuffle_ps(c, a, _MM_SHUFFLE(1, 1, 0, 0));
    __m128 v0 = _mm_shuffle_ps(u0, u1, _MM_SHUFFLE(2, 0, 2, 0));
    __m128 u2 = _mm_shuffle_ps(b, c, _MM_SHUFFLE(1, 1, 1, 1));
    __m128 u3 = _mm_shuffle_ps(a, b, _MM_SHUFFLE(2, 2, 2, 2));
    __m128 v1 = _mm_shuffle_ps(u2, u3, _MM_SHUFFLE(2, 0, 2, 0));
    __m128 u4 = _mm_shuffle_ps(c, a, _MM_SHUFFLE(3, 3, 2, 2));
    __m128 u5 = _mm_shuffle_ps(b, c, _MM_SHUFFLE(3, 3, 3, 3));
    __m128 v2 = _mm_shuffle_ps(u4, u5, _MM_SHUFFLE(2, 0, 2, 0));

    _mm_storeu_ps(ptr, v0);
    _mm_storeu_ps(ptr + 4, v1);
    _mm_storeu_ps(ptr + 8, v2);
}

//------------------------------------------------------------------------
//
// Blob-copy primitives namually vectored for SSE 4.2 (w/o OpenMP threads)
//
//------------------------------------------------------------------------

void blob_copy_4d_split_u8c3(const uint8_t* src_ptr, uint8_t* dst_ptr, size_t N_src_stride, size_t H_src_stride,
                             size_t N_dst_stride, size_t H_dst_stride, size_t C_dst_stride, int N, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int h = 0; h < H; h++) {
            const uint8_t* src = src_ptr + n * N_src_stride + h * H_src_stride;
            uint8_t* dst0 = dst_ptr + n * N_dst_stride + 0 * C_dst_stride + h * H_dst_stride;
            uint8_t* dst1 = dst_ptr + n * N_dst_stride + 1 * C_dst_stride + h * H_dst_stride;
            uint8_t* dst2 = dst_ptr + n * N_dst_stride + 2 * C_dst_stride + h * H_dst_stride;

            int w = 0;

            // SIMD128 manually
            for (; w < W - 16; w += 16) {
                __m128i r0, r1, r2;
                mm_load_deinterleave(&src[3 * w], r0, r1, r2);
                _mm_storeu_si128(reinterpret_cast<__m128i*>(dst0 + w), r0);
                _mm_storeu_si128(reinterpret_cast<__m128i*>(dst1 + w), r1);
                _mm_storeu_si128(reinterpret_cast<__m128i*>(dst2 + w), r2);
            }

            for (; w < W; w++) {
                dst0[w] = src[3 * w + 0];
                dst1[w] = src[3 * w + 1];
                dst2[w] = src[3 * w + 2];
            }
        }
}

void blob_copy_4d_split_f32c3(const float* src_ptr, float* dst_ptr, size_t N_src_stride, size_t H_src_stride,
                              size_t N_dst_stride, size_t H_dst_stride, size_t C_dst_stride, int N, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int h = 0; h < H; h++) {
            const float* src = src_ptr + n * N_src_stride + h * H_src_stride;
            float* dst0 = dst_ptr + n * N_dst_stride + 0 * C_dst_stride + h * H_dst_stride;
            float* dst1 = dst_ptr + n * N_dst_stride + 1 * C_dst_stride + h * H_dst_stride;
            float* dst2 = dst_ptr + n * N_dst_stride + 2 * C_dst_stride + h * H_dst_stride;

            int w = 0;

            // SIMD128 manually
            for (; w < W - 4; w += 4) {
                __m128 r0, r1, r2;
                mm_load_deinterleave(&src[3 * w], r0, r1, r2);
                _mm_storeu_ps(&dst0[w], r0);
                _mm_storeu_ps(&dst1[w], r1);
                _mm_storeu_ps(&dst2[w], r2);
            }

            for (; w < W; w++) {
                dst0[w] = src[3 * w + 0];
                dst1[w] = src[3 * w + 1];
                dst2[w] = src[3 * w + 2];
            }
        }
}

void blob_copy_4d_merge_u8c3(const uint8_t* src_ptr, uint8_t* dst_ptr, size_t N_src_stride, size_t H_src_stride,
                             size_t C_src_stride, size_t N_dst_stride, size_t H_dst_stride, int N, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int h = 0; h < H; h++) {
            const uint8_t* src0 = src_ptr + n * N_src_stride + 0 * C_src_stride + h * H_src_stride;
            const uint8_t* src1 = src_ptr + n * N_src_stride + 1 * C_src_stride + h * H_src_stride;
            const uint8_t* src2 = src_ptr + n * N_src_stride + 2 * C_src_stride + h * H_src_stride;

            uint8_t* dst = dst_ptr + n * N_dst_stride + h * H_dst_stride;

            int w = 0;

            // SIMD128 manually
            for (; w < W - 16; w += 16) {
                __m128i r0, r1, r2;
                r0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(src0 + w));
                r1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(src1 + w));
                r2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(src2 + w));
                mm_store_interleave(&dst[3 * w], r0, r1, r2);
            }

            for (; w < W; w++) {
                dst[3 * w + 0] = src0[w];
                dst[3 * w + 1] = src1[w];
                dst[3 * w + 2] = src2[w];
            }
        }
}

void blob_copy_4d_merge_f32c3(const float* src_ptr, float* dst_ptr, size_t N_src_stride, size_t H_src_stride,
                              size_t C_src_stride, size_t N_dst_stride, size_t H_dst_stride, int N, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int h = 0; h < H; h++) {
            const float* src0 = src_ptr + n * N_src_stride + 0 * C_src_stride + h * H_src_stride;
            const float* src1 = src_ptr + n * N_src_stride + 1 * C_src_stride + h * H_src_stride;
            const float* src2 = src_ptr + n * N_src_stride + 2 * C_src_stride + h * H_src_stride;

            float* dst = dst_ptr + n * N_dst_stride + h * H_dst_stride;

            int w = 0;

            // SIMD128 manually
            for (; w < W - 4; w += 4) {
                __m128 r0, r1, r2;
                r0 = _mm_loadu_ps(&src0[w]);
                r1 = _mm_loadu_ps(&src1[w]);
                r2 = _mm_loadu_ps(&src2[w]);
                mm_store_interleave(&dst[3 * w], r0, r1, r2);
            }

            for (; w < W; w++) {
                dst[3 * w + 0] = src0[w];
                dst[3 * w + 1] = src1[w];
                dst[3 * w + 2] = src2[w];
            }
        }
}

void blob_copy_5d_split_u8c3(const uint8_t* src_ptr, uint8_t* dst_ptr, size_t N_src_stride, size_t D_src_stride,
                             size_t H_src_stride, size_t N_dst_stride, size_t D_dst_stride, size_t H_dst_stride,
                             size_t C_dst_stride, int N, int D, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int d = 0; d < D; d++) {
            for (int h = 0; h < H; h++) {
                const uint8_t* src = src_ptr + n * N_src_stride + d * D_src_stride + h * H_src_stride;
                uint8_t* dst0 = dst_ptr + n * N_dst_stride + d * D_dst_stride + 0 * C_dst_stride + h * H_dst_stride;
                uint8_t* dst1 = dst_ptr + n * N_dst_stride + d * D_dst_stride + 1 * C_dst_stride + h * H_dst_stride;
                uint8_t* dst2 = dst_ptr + n * N_dst_stride + d * D_dst_stride + 2 * C_dst_stride + h * H_dst_stride;

                int w = 0;

                // SIMD128 manually
                for (; w < W - 16; w += 16) {
                    __m128i r0, r1, r2;
                    mm_load_deinterleave(&src[3 * w], r0, r1, r2);
                    _mm_storeu_si128(reinterpret_cast<__m128i*>(dst0 + w), r0);
                    _mm_storeu_si128(reinterpret_cast<__m128i*>(dst1 + w), r1);
                    _mm_storeu_si128(reinterpret_cast<__m128i*>(dst2 + w), r2);
                }

                for (; w < W; w++) {
                    dst0[w] = src[3 * w + 0];
                    dst1[w] = src[3 * w + 1];
                    dst2[w] = src[3 * w + 2];
                }
            }
        }
}

void blob_copy_5d_split_f32c3(const float* src_ptr, float* dst_ptr, size_t N_src_stride, size_t D_src_stride,
                              size_t H_src_stride, size_t N_dst_stride, size_t D_dst_stride, size_t H_dst_stride,
                              size_t C_dst_stride, int N, int D, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int d = 0; d < D; d++) {
            for (int h = 0; h < H; h++) {
                const float* src = src_ptr + n * N_src_stride + d * D_src_stride + h * H_src_stride;
                float* dst0 = dst_ptr + n * N_dst_stride + d * D_dst_stride + 0 * C_dst_stride + h * H_dst_stride;
                float* dst1 = dst_ptr + n * N_dst_stride + d * D_dst_stride + 1 * C_dst_stride + h * H_dst_stride;
                float* dst2 = dst_ptr + n * N_dst_stride + d * D_dst_stride + 2 * C_dst_stride + h * H_dst_stride;

                int w = 0;

                // SIMD128 manually
                for (; w < W - 4; w += 4) {
                    __m128 r0, r1, r2;
                    mm_load_deinterleave(&src[3 * w], r0, r1, r2);
                    _mm_storeu_ps(&dst0[w], r0);
                    _mm_storeu_ps(&dst1[w], r1);
                    _mm_storeu_ps(&dst2[w], r2);
                }

                for (; w < W; w++) {
                    dst0[w] = src[3 * w + 0];
                    dst1[w] = src[3 * w + 1];
                    dst2[w] = src[3 * w + 2];
                }
            }
        }
}

void blob_copy_5d_merge_u8c3(const uint8_t* src_ptr, uint8_t* dst_ptr, size_t N_src_stride, size_t D_src_stride,
                             size_t H_src_stride, size_t C_src_stride, size_t N_dst_stride, size_t D_dst_stride,
                             size_t H_dst_stride, int N, int D, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int d = 0; d < D; d++) {
            for (int h = 0; h < H; h++) {
                const uint8_t* src0 =
                    src_ptr + n * N_src_stride + 0 * C_src_stride + d * D_src_stride + h * H_src_stride;
                const uint8_t* src1 =
                    src_ptr + n * N_src_stride + 1 * C_src_stride + d * D_src_stride + h * H_src_stride;
                const uint8_t* src2 =
                    src_ptr + n * N_src_stride + 2 * C_src_stride + d * D_src_stride + h * H_src_stride;

                uint8_t* dst = dst_ptr + n * N_dst_stride + d * D_dst_stride + h * H_dst_stride;

                int w = 0;

                // SIMD128 manually
                for (; w < W - 16; w += 16) {
                    __m128i r0, r1, r2;
                    r0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(src0 + w));
                    r1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(src1 + w));
                    r2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(src2 + w));
                    mm_store_interleave(&dst[3 * w], r0, r1, r2);
                }

                for (; w < W; w++) {
                    dst[3 * w + 0] = src0[w];
                    dst[3 * w + 1] = src1[w];
                    dst[3 * w + 2] = src2[w];
                }
            }
        }
}

void blob_copy_5d_merge_f32c3(const float* src_ptr, float* dst_ptr, size_t N_src_stride, size_t D_src_stride,
                              size_t H_src_stride, size_t C_src_stride, size_t N_dst_stride, size_t D_dst_stride,
                              size_t H_dst_stride, int N, int D, int H, int W) {
    for (int n = 0; n < N; n++)
        for (int d = 0; d < D; d++) {
            for (int h = 0; h < H; h++) {
                const float* src0 = src_ptr + n * N_src_stride + 0 * C_src_stride + d * D_src_stride + h * H_src_stride;
                const float* src1 = src_ptr + n * N_src_stride + 1 * C_src_stride + d * D_src_stride + h * H_src_stride;
                const float* src2 = src_ptr + n * N_src_stride + 2 * C_src_stride + d * D_src_stride + h * H_src_stride;

                float* dst = dst_ptr + n * N_dst_stride + d * D_dst_stride + h * H_dst_stride;

                int w = 0;

                // SIMD128 manually
                for (; w < W - 4; w += 4) {
                    __m128 r0, r1, r2;
                    r0 = _mm_loadu_ps(&src0[w]);
                    r1 = _mm_loadu_ps(&src1[w]);
                    r2 = _mm_loadu_ps(&src2[w]);
                    mm_store_interleave(&dst[3 * w], r0, r1, r2);
                }

                for (; w < W; w++) {
                    dst[3 * w + 0] = src0[w];
                    dst[3 * w + 1] = src1[w];
                    dst[3 * w + 2] = src2[w];
                }
            }
        }
}

}  // namespace InferenceEngine
