// Copyright (c) Facebook, Inc. and its affiliates.
// All rights reserved.
//
// Copyright 2019 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

#pragma once

#include <stdint.h>

#include <xnnpack/common.h>
#include <xnnpack/scalar-utils.h>

#if XNN_ARCH_X86 || XNN_ARCH_X86_64
  #include <immintrin.h>
#endif

// The code below is adapted from Google's gemmlowp library.
// It is only used in XNNPACK unit tests and comparative benchmarks, but not the library itself.
//
// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


inline static int32_t gemmlowp_scalar_vqrdmulh_s32(int32_t a, int32_t b)
{
  const bool overflow = a == b && a == INT32_MIN;
  const int64_t ab_64 = (int64_t) a * (int64_t) b;
  const int32_t nudge = (a ^ b) >= 0 ? INT32_C(0x40000000) : -INT32_C(0x3FFFFFFF);
  const int32_t ab_x2_high32 = (int32_t) ((ab_64 + nudge) / INT64_C(0x80000000));
  return overflow ? INT32_MAX : ab_x2_high32;
}

inline static int32_t gemmlowp_scalar_rdivbypo2_s32(int32_t x, int exponent)
{
  const int32_t mask = ((1 << exponent) - 1);
  const int32_t remainder = x & mask;
  const int32_t threshold = (mask >> 1) + (int32_t) (x < 0);
  return asr_s32(x, exponent) + (int32_t) (remainder > threshold);
}

#if XNN_ARCH_X86 || XNN_ARCH_X86_64
  static inline __m128i gemmlowp_sse_rdivbypo2_s32(__m128i x, int exponent) {
    const __m128i mask = _mm_set1_epi32((int32_t) ((UINT64_C(1) << exponent) - UINT64_C(1)));
    const __m128i remainder = _mm_and_si128(x, mask);
    const __m128i threshold = _mm_sub_epi32(
        _mm_srli_epi32(mask, 1), _mm_cmplt_epi32(x, _mm_setzero_si128()));
    return _mm_sub_epi32(
        _mm_sra_epi32(x, _mm_cvtsi32_si128(exponent)),
        _mm_cmpgt_epi32(remainder, threshold));
  }

  static inline __m128i gemmlowp_sse_mul_s32(__m128i a, __m128i b) {
  #ifdef __SSE4_1__
    return _mm_mul_epi32(a, b);
  #else
    __m128i sign, zero, mul_us, a_neg, b_neg, mul_us_neg;
    sign = _mm_xor_si128(a, b);
    sign = _mm_srai_epi32(sign, 31); // promote sign bit to all fields, all fff if
                                     // negative and all 0 if positive
    sign = _mm_shuffle_epi32(
        sign,
        _MM_SHUFFLE(2, 2, 0, 0)); // promote sign bit to 3 and 1st data lanes
    zero = _mm_setzero_si128();
  #ifdef __SSSE3__
    a_neg = _mm_abs_epi32(a); // negate a and b
    b_neg = _mm_abs_epi32(b); // negate a and b
  #else  // pre-SSSE3
    const __m128i a_neg_mask = _mm_cmplt_epi32(a, zero);
    a_neg = _mm_sub_epi32(_mm_xor_si128(a, a_neg_mask), a_neg_mask);
    const __m128i b_neg_mask = _mm_cmplt_epi32(b, zero);
    b_neg = _mm_sub_epi32(_mm_xor_si128(b, b_neg_mask), b_neg_mask);
  #endif  // pre-SSSE3
    mul_us = _mm_mul_epu32(a_neg, b_neg); // uses 0 and 2nd data lanes, (abs), the
                                          // multiplication gives 64 bit result
    mul_us_neg = _mm_sub_epi64(zero, mul_us);
    mul_us_neg = _mm_and_si128(sign, mul_us_neg);
    mul_us = _mm_andnot_si128(sign, mul_us);
    return _mm_or_si128(mul_us, mul_us_neg);
  #endif
  }

  static inline __m128i gemmlowp_sse_vqrdmulh_s32(__m128i a, __m128i b) {
    // saturation only happen if a == b == INT32_MIN
    const __m128i min = _mm_set1_epi32(INT32_MIN);
    const __m128i saturation_mask =
        _mm_and_si128(_mm_cmpeq_epi32(a, b), _mm_cmpeq_epi32(a, min));

    // a = a0 | a1 | a2 | a3
    // b = b0 | b1 | b2 | b3
    const __m128i a0_a2 = a;
    const __m128i a1_a3 = _mm_srli_si128(a, 4);
    const __m128i b0_b2 = b;
    const __m128i b1_b3 = _mm_srli_si128(b, 4);

    const __m128i a0b0_a2b2 = gemmlowp_sse_mul_s32(a0_a2, b0_b2);
    const __m128i a1b1_a3b3 = gemmlowp_sse_mul_s32(a1_a3, b1_b3);

    // do the rounding and take into account that it will be doubled
    const __m128i nudge = _mm_set1_epi64x(1 << 30);
    const __m128i a0b0_a2b2_rounded = _mm_add_epi64(a0b0_a2b2, nudge);
    const __m128i a1b1_a3b3_rounded = _mm_add_epi64(a1b1_a3b3, nudge);

    // do the doubling
    const __m128i a0b0_a2b2_rounded_2x = _mm_slli_epi64(a0b0_a2b2_rounded, 1);
    const __m128i a1b1_a3b3_rounded_2x = _mm_slli_epi64(a1b1_a3b3_rounded, 1);

  // get the high part of the products
  #ifdef __SSE4_1__
    const __m128i result = _mm_blend_epi16(
        _mm_srli_epi64(a0b0_a2b2_rounded_2x, 32), a1b1_a3b3_rounded_2x, 0xCC);
  #else
    const __m128i result0213 = _mm_castps_si128(_mm_shuffle_ps(
        _mm_castsi128_ps(a0b0_a2b2_rounded_2x),
        _mm_castsi128_ps(a1b1_a3b3_rounded_2x),
        _MM_SHUFFLE(3, 1, 3, 1)));
    const __m128i result = _mm_shuffle_epi32(result0213, _MM_SHUFFLE(3, 1, 2, 0));
  #endif

  // saturate those which overflowed
  #ifdef __SSE4_1__
    const __m128i saturated_result = _mm_blendv_epi8(result, min, saturation_mask);
  #else
    const __m128i saturated_result = _mm_or_si128(
        _mm_and_si128(saturation_mask, min),
        _mm_andnot_si128(saturation_mask, result));
  #endif
    return saturated_result;
  }
#endif  // XNN_ARCH_X86 || XNN_ARCH_X86_64
