/* Copyright 2016 The TensorFlow Authors. 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.
==============================================================================*/

#ifndef TENSORFLOW_CORE_KERNELS_SPARSE_MATMUL_OP_H_
#define TENSORFLOW_CORE_KERNELS_SPARSE_MATMUL_OP_H_

#include "third_party/eigen3/Eigen/Core"
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/types.h"

#if defined(PLATFORM_WINDOWS)
#include "tensorflow/core/platform/windows/cpu_info.h"
#include "tensorflow/core/platform/windows/intrinsics_port.h"
#endif

namespace Eigen {
namespace internal {

// Return the float representation of the bfloat16 value
// in the lower 16-bits of input
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pexpand_bf16_l(const Packet& from) {
  tensorflow::uint32 tmp;
#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
  tmp = (reinterpret_cast<const tensorflow::uint32&>(from)) & 0xffff0000;
#else
  tmp = (reinterpret_cast<const tensorflow::uint32&>(from) << 16) & 0xffff0000;
#endif
  return reinterpret_cast<const float&>(tmp);
}

// Return the float representation of the bfloat16 value
// in the upper 16-bits of input
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pexpand_bf16_u(const Packet& from) {
  tensorflow::uint32 tmp;
#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
  tmp = (reinterpret_cast<const tensorflow::uint32&>(from) << 16) & 0xffff0000;
#else
  tmp = (reinterpret_cast<const tensorflow::uint32&>(from)) & 0xffff0000;
#endif
  return reinterpret_cast<const float&>(tmp);
}

// Specialization non-scalar version on non-sse.
// Enable vectorization on z13 and higher
#if defined(EIGEN_VECTORIZE_ALTIVEC) || defined(EIGEN_VECTORIZE_VSX) || \
    defined(EIGEN_VECTORIZE_NEON) || defined(EIGEN_VECTORIZE_ZVECTOR)
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet4f pexpand_bf16_l(const Packet4f& from) {
  float r[4];
  tensorflow::uint32 p[4];
  pstoreu(r, from);
  tensorflow::uint32* ir = reinterpret_cast<tensorflow::uint32*>(r);
  p[0] = (ir[0] << 16) & 0xffff0000;
  p[1] = ir[0] & 0xffff0000;
  p[2] = (ir[1] << 16) & 0xffff0000;
  p[3] = ir[1] & 0xffff0000;
  return ploadu<Packet4f>(reinterpret_cast<float*>(p));
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet4f pexpand_bf16_u(const Packet4f& from) {
  float r[4];
  tensorflow::uint32 p[4];
  pstoreu(r, from);
  tensorflow::uint32* ir = reinterpret_cast<tensorflow::uint32*>(r);
  p[0] = (ir[2] << 16) & 0xffff0000;
  p[1] = ir[2] & 0xffff0000;
  p[2] = (ir[3] << 16) & 0xffff0000;
  p[3] = ir[3] & 0xffff0000;
  return ploadu<Packet4f>(reinterpret_cast<float*>(p));
}
#endif

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pinterleave4x64(const Packet& from) {
  return from;
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pbroadcast_first(const Packet& a) {
  return a;
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pbroadcast_second(const Packet& a) {
  assert(false && "Not applicable to Scalar Values");
  return a;
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pbroadcast_third(const Packet& a) {
  assert(false && "Not applicable to Scalar Values");
  return a;
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pbroadcast_fourth(const Packet& a) {
  assert(false && "Not applicable to Scalar Values");
  return a;
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pload4bf16(
    const typename unpacket_traits<Packet>::type* from) {
  assert(false && "Not applicable to Scalar Values");
  return Packet();
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet pload2bf16(
    const typename unpacket_traits<Packet>::type* from) {
  assert(false && "Not applicable to Scalar Values");
  return Packet();
}

// Specialization for pload4bf16 and pload2bf16 for non-sse.
// Enable vectorization on z13 and higher.
#if defined(EIGEN_VECTORIZE_ALTIVEC) || defined(EIGEN_VECTORIZE_VSX) || \
    defined(EIGEN_VECTORIZE_NEON) || defined(EIGEN_VECTORIZE_ZVECTOR)
template <>
EIGEN_STRONG_INLINE Packet4f pload4bf16<Packet4f>(const float* from) {
  tensorflow::uint32 p[4];
  const tensorflow::uint32* ir =
      reinterpret_cast<const tensorflow::uint32*>(from);
  p[0] = (ir[0] << 16) & 0xffff0000;
  p[1] = ir[0] & 0xffff0000;
  p[2] = (ir[1] << 16) & 0xffff0000;
  p[3] = ir[1] & 0xffff0000;
  return ploadu<Packet4f>(reinterpret_cast<float*>(p));
}

template <>
EIGEN_STRONG_INLINE Packet4f pload2bf16<Packet4f>(const float* from) {
  tensorflow::uint32 p[4];
  const tensorflow::uint32* ir =
      reinterpret_cast<const tensorflow::uint32*>(from);
  p[0] = (ir[0] << 16) & 0xffff0000;
  p[1] = ir[0] & 0xffff0000;
  p[2] = (ir[0] << 16) & 0xffff0000;
  p[3] = ir[0] & 0xffff0000;
  return ploadu<Packet4f>(reinterpret_cast<float*>(p));
}
#endif

#if defined(EIGEN_VECTORIZE_ALTIVEC) || defined(EIGEN_VECTORIZE_VSX)
// Return a packet with the first value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_first<Packet4f>(const Packet4f& a) {
  return vec_splat(a, 0);
}

// Return a packet with the second value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_second<Packet4f>(const Packet4f& a) {
  return vec_splat(a, 1);
}

// Return a packet with the third value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_third<Packet4f>(const Packet4f& a) {
  return vec_splat(a, 2);
}

// Return a packet with the fourth value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_fourth<Packet4f>(const Packet4f& a) {
  return vec_splat(a, 3);
}
#endif

#ifdef EIGEN_VECTORIZE_SSE2
// For PacketSize of 4 floats the Packet is not modified
template <>
EIGEN_STRONG_INLINE Packet4f pinterleave4x64<Packet4f>(const Packet4f& from) {
  return from;
}

// Return a Packet with 4 floats loaded from 4 bfloat16 values
template <>
EIGEN_STRONG_INLINE Packet4f pload4bf16<Packet4f>(const float* from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castpd_si128(_mm_load_pd1((const double*)from));
  return _mm_castsi128_ps(_mm_unpacklo_epi16(zero, tmp));
}

// Return a Packet with 2 floats loaded from 2 bfloat16 values
template <>
EIGEN_STRONG_INLINE Packet4f pload2bf16<Packet4f>(const float* from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castps_si128(_mm_load_ps1(from));
  return _mm_castsi128_ps(_mm_unpacklo_epi16(zero, tmp));
}

// Return a Packet with 4 floats expanded from 4 bfloat16 values
// in the lower half of the 128-bit lane
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet4f pexpand_bf16_l(const Packet4f& from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castps_si128(from);
  return _mm_castsi128_ps(_mm_unpacklo_epi16(zero, tmp));
}

// Return a Packet with 4 floats expanded from 4 bfloat16 values
// in the upper half of the 128-bit lane
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet4f pexpand_bf16_u(const Packet4f& from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castps_si128(from);
  return _mm_castsi128_ps(_mm_unpackhi_epi16(zero, tmp));
}

// Return a packet with the first value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_first<Packet4f>(const Packet4f& a) {
  return _mm_set1_ps(pfirst<Packet4f>(a));
}

// Return a packet with the second value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_second<Packet4f>(const Packet4f& a) {
  return _mm_set1_ps(_mm_cvtss_f32(_mm_shuffle_ps(a, a, 1)));
}

// Return a packet with the third value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_third<Packet4f>(const Packet4f& a) {
  return _mm_set1_ps(_mm_cvtss_f32(_mm_shuffle_ps(a, a, 2)));
}

// Return a packet with the fourth value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet4f pbroadcast_fourth<Packet4f>(const Packet4f& a) {
  return _mm_set1_ps(_mm_cvtss_f32(_mm_shuffle_ps(a, a, 3)));
}

#endif

#ifdef EIGEN_VECTORIZE_AVX512
template <>
EIGEN_STRONG_INLINE Packet16f
pbroadcast_first<Packet16f>(const Packet16f& a_in) {
  Packet4f a = _mm512_castps512_ps128(a_in);
  return _mm512_broadcastss_ps(a);
}
template <>
EIGEN_STRONG_INLINE Packet16f
pbroadcast_second<Packet16f>(const Packet16f& a_in) {
  Packet4f a = _mm512_castps512_ps128(a_in);
  return _mm512_broadcastss_ps(_mm_shuffle_ps(a, a, _MM_SHUFFLE(1, 1, 1, 1)));
}
template <>
EIGEN_STRONG_INLINE Packet16f
pbroadcast_third<Packet16f>(const Packet16f& a_in) {
  Packet4f a = _mm512_castps512_ps128(a_in);
  return _mm512_broadcastss_ps(_mm_shuffle_ps(a, a, _MM_SHUFFLE(2, 2, 2, 2)));
}
template <>
EIGEN_STRONG_INLINE Packet16f
pbroadcast_fourth<Packet16f>(const Packet16f& a_in) {
  Packet4f a = _mm512_castps512_ps128(a_in);
  return _mm512_broadcastss_ps(_mm_shuffle_ps(a, a, _MM_SHUFFLE(3, 3, 3, 3)));
}
template <>
EIGEN_STRONG_INLINE Packet8d pbroadcast_first<Packet8d>(const Packet8d& a_in) {
  Packet2d a = _mm512_castpd512_pd128(a_in);
  return _mm512_broadcastsd_pd(a);
}
template <>
EIGEN_STRONG_INLINE Packet8d pbroadcast_second<Packet8d>(const Packet8d& a_in) {
  Packet2d a = _mm_permute_pd(_mm512_castpd512_pd128(a_in), 3);
  return _mm512_broadcastsd_pd(a);
}
template <>
EIGEN_STRONG_INLINE Packet8d pbroadcast_third<Packet8d>(const Packet8d& a_in) {
  Packet2d a = _mm256_extractf128_pd(_mm512_castpd512_pd256(a_in), 1);
  return _mm512_broadcastsd_pd(a);
}
template <>
EIGEN_STRONG_INLINE Packet8d pbroadcast_fourth<Packet8d>(const Packet8d& a_in) {
  Packet2d a =
      _mm_permute_pd(_mm256_extractf128_pd(_mm512_castpd512_pd256(a_in), 1), 3);
  return _mm512_broadcastsd_pd(a);
}
template <>
EIGEN_STRONG_INLINE Packet16i
pbroadcast_first<Packet16i>(const Packet16i& a_in) {
  Packet4i a = _mm512_castsi512_si128(a_in);
  return _mm512_broadcastd_epi32(a);
}
template <>
EIGEN_STRONG_INLINE Packet16i
pbroadcast_second<Packet16i>(const Packet16i& a_in) {
  Packet4i a = _mm512_castsi512_si128(a_in);
  return _mm512_broadcastd_epi32(_mm_shuffle_epi32(a, _MM_SHUFFLE(1, 1, 1, 1)));
}
template <>
EIGEN_STRONG_INLINE Packet16i
pbroadcast_third<Packet16i>(const Packet16i& a_in) {
  Packet4i a = _mm512_castsi512_si128(a_in);
  return _mm512_broadcastd_epi32(_mm_shuffle_epi32(a, _MM_SHUFFLE(2, 2, 2, 2)));
}
template <>
EIGEN_STRONG_INLINE Packet16i
pbroadcast_fourth<Packet16i>(const Packet16i& a_in) {
  Packet4i a = _mm512_castsi512_si128(a_in);
  return _mm512_broadcastd_epi32(_mm_shuffle_epi32(a, _MM_SHUFFLE(3, 3, 3, 3)));
}
#endif

#ifdef EIGEN_VECTORIZE_AVX
// For a Packet of Size 8 floats(256-bits), swap the 2nd and 3rd quadwords
template <>
EIGEN_STRONG_INLINE Packet8f pinterleave4x64<Packet8f>(const Packet8f& from) {
#ifdef EIGEN_VECTORIZE_AVX2
  return _mm256_castsi256_ps(_mm256_permute4x64_epi64(_mm256_castps_si256(from),
                                                      _MM_SHUFFLE(3, 1, 2, 0)));
#else
  auto tmp1 = _mm256_extract_epi32(_mm256_castps_si256(from), 2);
  auto tmp2 = _mm256_extract_epi32(_mm256_castps_si256(from), 3);
  auto tmp3 = _mm256_extract_epi32(_mm256_castps_si256(from), 4);
  auto tmp4 = _mm256_extract_epi32(_mm256_castps_si256(from), 5);
  auto tmp5 = _mm256_insert_epi32(_mm256_castps_si256(from), tmp1, 4);
  tmp5 = _mm256_insert_epi32(tmp5, tmp2, 5);
  tmp5 = _mm256_insert_epi32(tmp5, tmp3, 2);
  tmp5 = _mm256_insert_epi32(tmp5, tmp4, 3);
  return _mm256_castsi256_ps(tmp5);
#endif
}
// Return a Packet with 4 floats loaded from 4 bfloat16 values
template <>
EIGEN_STRONG_INLINE Packet8f pload4bf16<Packet8f>(const float* from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castpd_si128(_mm_load_pd1((const double*)from));
  return _mm256_castps128_ps256(
      _mm_castsi128_ps(_mm_unpacklo_epi16(zero, tmp)));
}
// Return a Packet with 2 floats loaded from 2 bfloat16 values
template <>
EIGEN_STRONG_INLINE Packet8f pload2bf16<Packet8f>(const float* from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castps_si128(_mm_load_ps1(from));
  return _mm256_castps128_ps256(
      _mm_castsi128_ps(_mm_unpacklo_epi16(zero, tmp)));
}

#ifdef EIGEN_VECTORIZE_AVX512
// Return a Packet with 4 floats loaded from 4 bfloat16 values
template <>
EIGEN_STRONG_INLINE Packet16f pload4bf16<Packet16f>(const float* from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castpd_si128(_mm_load_pd1((const double*)from));
  return _mm512_castps128_ps512(
      _mm_castsi128_ps(_mm_unpacklo_epi16(zero, tmp)));
}
// Return a Packet with 2 floats loaded from 2 bfloat16 values
template <>
EIGEN_STRONG_INLINE Packet16f pload2bf16<Packet16f>(const float* from) {
  __m128i zero = _mm_setzero_si128();
  __m128i tmp = _mm_castps_si128(_mm_load_ps1(from));
  return _mm512_castps128_ps512(
      _mm_castsi128_ps(_mm_unpacklo_epi16(zero, tmp)));
}
#endif

// For each 128-bit lane convert 4 bfloat to 4 float values from the lower half
// of the 128-bit lane
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet8f pexpand_bf16_l(const Packet8f& from) {
#ifdef EIGEN_VECTORIZE_AVX2
  __m256i zero = _mm256_setzero_si256();
  __m256i tmp = _mm256_castps_si256(from);
  return _mm256_castsi256_ps(_mm256_unpacklo_epi16(zero, tmp));
#else
  __m128i zero = _mm_setzero_si128();
  __m128i low = _mm_castps_si128(_mm256_extractf128_ps(from, 0));
  __m128i res_l = _mm_unpacklo_epi16(zero, low);
  __m128i high = _mm_castps_si128(_mm256_extractf128_ps(from, 1));
  __m128i res_h = _mm_unpacklo_epi16(zero, high);
  __m256 res = _mm256_castps128_ps256(_mm_castsi128_ps(res_l));
  res = _mm256_insertf128_ps(res, _mm_castsi128_ps(res_h), 1);
  return res;
#endif
}

// For each 128-bit lane convert 4 bfloat to 4 float values from the upper half
// of the 128-bit lane
template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet8f pexpand_bf16_u(const Packet8f& from) {
#ifdef EIGEN_VECTORIZE_AVX2
  __m256i zero = _mm256_setzero_si256();
  __m256i tmp = _mm256_castps_si256(from);
  return _mm256_castsi256_ps(_mm256_unpackhi_epi16(zero, tmp));
#else
  __m128i zero = _mm_setzero_si128();
  __m128i low = _mm_castps_si128(_mm256_extractf128_ps(from, 0));
  __m128i res_l = _mm_unpackhi_epi16(zero, low);
  __m128i high = _mm_castps_si128(_mm256_extractf128_ps(from, 1));
  __m128i res_h = _mm_unpackhi_epi16(zero, high);
  __m256 res = _mm256_castps128_ps256(_mm_castsi128_ps(res_l));
  res = _mm256_insertf128_ps(res, _mm_castsi128_ps(res_h), 1);
  return res;
#endif
}

// Return a packet with the first value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet8f pbroadcast_first<Packet8f>(const Packet8f& a) {
  return _mm256_set1_ps(pfirst<Packet8f>(a));
}

// Return a packet with the second value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet8f pbroadcast_second<Packet8f>(const Packet8f& a) {
  return _mm256_set1_ps(
      _mm_cvtss_f32(_mm256_castps256_ps128(_mm256_permute_ps(a, 1))));
}

// Return a packet with the third value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet8f pbroadcast_third<Packet8f>(const Packet8f& a) {
  return _mm256_set1_ps(
      _mm_cvtss_f32(_mm256_castps256_ps128(_mm256_permute_ps(a, 2))));
}

// Return a packet with the fourth value of the input Packet replicated
template <>
EIGEN_STRONG_INLINE Packet8f pbroadcast_fourth<Packet8f>(const Packet8f& a) {
  return _mm256_set1_ps(
      _mm_cvtss_f32(_mm256_castps256_ps128(_mm256_permute_ps(a, 3))));
}

#endif

#ifdef EIGEN_VECTORIZE_AVX512

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet16f pexpand_bf16_l(const Packet16f& from) {
  return _mm512_castsi512_ps(_mm512_slli_epi32(
      _mm512_cvtepu16_epi32(_mm512_castsi512_si256(_mm512_castps_si512(from))),
      16));
}

template <typename Packet>
EIGEN_DEVICE_FUNC inline Packet16f pexpand_bf16_u(const Packet16f& from) {
  Packet16i tmp = _mm512_castps_si512(from);
  Packet16i tmp2 = _mm512_alignr_epi32(tmp, tmp, 8);
  return _mm512_castsi512_ps(_mm512_slli_epi32(
      _mm512_cvtepu16_epi32(_mm512_castsi512_si256(tmp2)), 16));
}

#endif
}  // namespace internal
}  // namespace Eigen
#endif  // TENSORFLOW_CORE_KERNELS_SPARSE_MATMUL_OP_H_