vec_ops.h

// Implementation of <simd> -*- C++ -*-
 
// Copyright The GNU Toolchain Authors.
//
// This file is part of the GNU ISO C++ Library.  This library is free
// software; you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the
// Free Software Foundation; either version 3, or (at your option)
// any later version.
 
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
 
// Under Section 7 of GPL version 3, you are granted additional
// permissions described in the GCC Runtime Library Exception, version
// 3.1, as published by the Free Software Foundation.
 
// You should have received a copy of the GNU General Public License and
// a copy of the GCC Runtime Library Exception along with this program;
// see the files COPYING3 and COPYING.RUNTIME respectively.  If not, see
// <http://www.gnu.org/licenses/>.
 
#ifndef _GLIBCXX_VEC_OPS_H
#define _GLIBCXX_VEC_OPS_H 1
 
#ifdef _GLIBCXX_SYSHDR
#pragma GCC system_header
#endif
 
#if __cplusplus >= 202400L
 
#include "simd_details.h"
 
#include <bit>
#include <bits/utility.h>
 
// psabi warnings are bogus because the ABI of the internal types never leaks into user code
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpsabi"
 
namespace std _GLIBCXX_VISIBILITY(default)
{
_GLIBCXX_BEGIN_NAMESPACE_VERSION
namespace simd
{
  template <std::signed_integral _Tp>
    constexpr bool
    __signed_has_single_bit(_Tp __x)
    { return __has_single_bit(make_unsigned_t<_Tp>(__x)); }
 
  /**
   * Alias for a vector builtin with given value type and total sizeof.
   */
  template <__vectorizable _Tp, size_t _Bytes>
    requires (__has_single_bit(_Bytes))
    using __vec_builtin_type_bytes [[__gnu__::__vector_size__(_Bytes)]] = _Tp;
 
  /**
   * Alias for a vector builtin with given value type @p _Tp and @p _Width.
   */
  template <__vectorizable _Tp, __simd_size_type _Width>
    requires (__signed_has_single_bit(_Width))
    using __vec_builtin_type = __vec_builtin_type_bytes<_Tp, sizeof(_Tp) * _Width>;
 
  /**
   * Constrain to any vector builtin with given value type and optional width.
   */
  template <typename _Tp, typename _ValueType,
            __simd_size_type _Width = sizeof(_Tp) / sizeof(_ValueType)>
    concept __vec_builtin_of
      = !is_class_v<_Tp> && !is_pointer_v<_Tp> && !is_arithmetic_v<_Tp>
          && __vectorizable<_ValueType>
          && _Width >= 1 && sizeof(_Tp) / sizeof(_ValueType) == _Width
          && same_as<__vec_builtin_type_bytes<_ValueType, sizeof(_Tp)>, _Tp>
          && requires(_Tp& __v, _ValueType __x) { __v[0] = __x; };
 
  /**
   * Constrain to any vector builtin.
   */
  template <typename _Tp>
    concept __vec_builtin
      = __vec_builtin_of<_Tp, remove_cvref_t<decltype(declval<const _Tp>()[0])>>;
 
  /**
   * Alias for the value type of the given __vec_builtin type @p _Tp.
   */
  template <__vec_builtin _Tp>
    using __vec_value_type = remove_cvref_t<decltype(declval<const _Tp>()[0])>;
 
  /**
   * The width (number of value_type elements) of the given vector builtin or arithmetic type.
   */
  template <typename _Tp>
    inline constexpr __simd_size_type __width_of = 1;
 
  template <typename _Tp>
    requires __vec_builtin<_Tp>
    inline constexpr __simd_size_type __width_of<_Tp> = sizeof(_Tp) / sizeof(__vec_value_type<_Tp>);
 
  /**
   * Alias for a vector builtin with equal value type and new width @p _Np.
   */
  template <__simd_size_type _Np, __vec_builtin _TV>
    using __resize_vec_builtin_t = __vec_builtin_type<__vec_value_type<_TV>, _Np>;
 
  template <__vec_builtin _TV>
    requires (__width_of<_TV> > 1)
    using __half_vec_builtin_t = __resize_vec_builtin_t<__width_of<_TV> / 2, _TV>;
 
  template <__vec_builtin _TV>
    using __double_vec_builtin_t = __resize_vec_builtin_t<__width_of<_TV> * 2, _TV>;
 
  template <typename _Up, __vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr __vec_builtin_type_bytes<_Up, sizeof(_TV)>
    __vec_bit_cast(_TV __v)
    { return reinterpret_cast<__vec_builtin_type_bytes<_Up, sizeof(_TV)>>(__v); }
 
  template <int _Np, __vec_builtin _TV>
    requires signed_integral<__vec_value_type<_TV>>
    static constexpr _TV _S_vec_implicit_mask = []<int... _Is> (integer_sequence<int, _Is...>) {
      return _TV{ (_Is < _Np ? -1 : 0)... };
    } (make_integer_sequence<int, __width_of<_TV>>());
 
  /**
   * Helper function to work around Clang not allowing v[i] in constant expressions.
   */
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr __vec_value_type<_TV>
    __vec_get(_TV __v, int __i)
    {
#ifdef _GLIBCXX_CLANG
      if consteval
        {
          return __builtin_bit_cast(array<__vec_value_type<_TV>, __width_of<_TV>>, __v)[__i];
        }
      else
#endif
        {
          return __v[__i];
        }
    }
 
  /**
   * Helper function to work around Clang and GCC not allowing assignment to v[i] in constant
   * expressions.
   */
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr void
    __vec_set(_TV& __v, int __i, __vec_value_type<_TV> __x)
    {
      if consteval
        {
#ifdef _GLIBCXX_CLANG
          auto __arr = __builtin_bit_cast(array<__vec_value_type<_TV>, __width_of<_TV>>, __v);
          __arr[__i] = __x;
          __v = __builtin_bit_cast(_TV, __arr);
#else
          constexpr auto [...__j] = _IotaArray<__width_of<_TV>>;
          __v = _TV{(__i == __j ? __x : __v[__j])...};
#endif
        }
      else
        {
          __v[__i] = __x;
        }
    }
 
  /** @internal
   * Return vector builtin with all values from @p __a and @p __b.
   */
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr __vec_builtin_type<__vec_value_type<_TV>, __width_of<_TV> * 2>
    __vec_concat(_TV __a, _TV __b)
    {
      constexpr auto [...__is] = _IotaArray<__width_of<_TV> * 2>;
      return __builtin_shufflevector(__a, __b, __is...);
    }
 
  /** @internal
   * Concatenate the first @p _N0 elements from @p __a with the first @p _N1 elements from @p __b
   * with the elements from applying this function recursively to @p __rest.
   *
   * @pre _N0 <= __width_of<_TV0> && _N1 <= __width_of<_TV1> && _Ns <= __width_of<_TVs> && ...
   *
   * Strategy: Aim for a power-of-2 tree concat. E.g.
   * - cat(2, 2, 2, 2) -> cat(4, 2, 2) -> cat(4, 4)
   * - cat(2, 2, 2, 2, 8) -> cat(4, 2, 2, 8) -> cat(4, 4, 8) -> cat(8, 8)
   */
  template <int _N0, int _N1, int... _Ns, __vec_builtin _TV0, __vec_builtin _TV1,
           __vec_builtin... _TVs>
    [[__gnu__::__always_inline__]]
    constexpr __vec_builtin_type<__vec_value_type<_TV0>,
                                 __bit_ceil(unsigned(_N0 + (_N1 + ... + _Ns)))>
    __vec_concat_sized(const _TV0& __a, const _TV1& __b, const _TVs&... __rest);
 
  template <int _N0, int _N1, int _N2, int... _Ns, __vec_builtin _TV0, __vec_builtin _TV1,
            __vec_builtin _TV2, __vec_builtin... _TVs>
    requires (__has_single_bit(unsigned(_N0))) && (_N0 >= (_N1 + _N2))
    [[__gnu__::__always_inline__]]
    constexpr __vec_builtin_type<__vec_value_type<_TV0>,
                                 __bit_ceil(unsigned(_N0 + _N1 + (_N2 + ... + _Ns)))>
    __vec_concat_sized(const _TV0& __a, const _TV1& __b, const _TV2& __c, const _TVs&... __rest)
    {
      return __vec_concat_sized<_N0, _N1 + _N2, _Ns...>(
               __a, __vec_concat_sized<_N1, _N2>(__b, __c), __rest...);
    }
 
  template <int _N0, int _N1, int... _Ns, __vec_builtin _TV0, __vec_builtin _TV1,
           __vec_builtin... _TVs>
    [[__gnu__::__always_inline__]]
    constexpr __vec_builtin_type<__vec_value_type<_TV0>,
                                 __bit_ceil(unsigned(_N0 + (_N1 + ... + _Ns)))>
    __vec_concat_sized(const _TV0& __a, const _TV1& __b, const _TVs&... __rest)
    {
      // __is is rounded up because we need to generate a power-of-2 vector:
      constexpr auto [...__is] = _IotaArray<__bit_ceil(unsigned(_N0 + _N1)), int>;
      const auto __ab = __builtin_shufflevector(__a, __b, [](int __i) consteval {
                          if (__i < _N0) // copy from __a
                            return __i;
                          else if (__i < _N0 + _N1) // copy from __b
                            return __i - _N0 + __width_of<_TV0>; // _N0 <= __width_of<_TV0>
                          else // can't index into __rest
                            return -1; // don't care
                        }(__is)...);
      if constexpr (sizeof...(__rest) == 0)
        return __ab;
      else
        return __vec_concat_sized<_N0 + _N1, _Ns...>(__ab, __rest...);
    }
 
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr __half_vec_builtin_t<_TV>
    __vec_split_lo(_TV __v)
    {
      constexpr int __n = __width_of<_TV> / 2;
      constexpr auto [...__is] = _IotaArray<__n>;
      return __builtin_shufflevector(__v, __v, __is...);
    }
 
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr __half_vec_builtin_t<_TV>
    __vec_split_hi(_TV __v)
    {
      constexpr int __n = __width_of<_TV> / 2;
      constexpr auto [...__is] = _IotaArray<__n>;
      return __builtin_shufflevector(__v, __v, (__n + __is)...);
    }
 
  /** @internal
   * Return @p __x zero-padded to @p _Bytes bytes.
   *
   * Use this function when you need two objects of the same size (e.g. for __vec_concat).
   */
  template <size_t _Bytes, __vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr auto
    __vec_zero_pad_to(_TV __x)
    {
      if constexpr (sizeof(_TV) == _Bytes)
        return __x;
      else if constexpr (sizeof(_TV) <= sizeof(0ull))
        {
          using _Up = _UInt<sizeof(_TV)>;
          __vec_builtin_type_bytes<_Up, _Bytes> __tmp = {__builtin_bit_cast(_Up, __x)};
          return __builtin_bit_cast(__vec_builtin_type_bytes<__vec_value_type<_TV>, _Bytes>, __tmp);
        }
      else if constexpr (sizeof(_TV) < _Bytes)
        return __vec_zero_pad_to<_Bytes>(__vec_concat(__x, _TV()));
      else
        static_assert(false);
    }
 
  /** @internal
   * Return a type with sizeof 16, add zero-padding to @p __x. The input must be smaller.
   *
   * Use this function instead of the above when you need to pad an argument for a SIMD builtin.
   */
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr auto
    __vec_zero_pad_to_16(_TV __x)
    {
      static_assert(sizeof(_TV) < 16);
      return __vec_zero_pad_to<16>(__x);
    }
 
  // work around __builtin_constant_p returning false unless passed a variable
  // (__builtin_constant_p(x[0]) is false while __is_const_known(x[0]) is true)
  template <typename _Tp>
    [[__gnu__::__always_inline__]]
    constexpr bool
    __is_const_known(const _Tp& __x)
    {
      return __builtin_constant_p(__x);
    }
 
  [[__gnu__::__always_inline__]]
  constexpr bool
  __is_const_known(const auto&... __xs) requires(sizeof...(__xs) >= 2)
  {
    if consteval
      {
        return true;
      }
    else
      {
        return (__is_const_known(__xs) && ...);
      }
  }
 
  [[__gnu__::__always_inline__]]
  constexpr bool
  __is_const_known_equal_to(const auto& __x, const auto& __expect)
  { return __is_const_known(__x == __expect) && __x == __expect; }
 
#if _GLIBCXX_X86
  template <__vec_builtin _UV, __vec_builtin _TV>
    inline _UV
    __x86_cvt_f16c(_TV __v);
#endif
 
 
  /** @internal
   * Simple wrapper around __builtin_convertvector to provide static_cast-like syntax.
   *
   * Works around GCC failing to use the F16C/AVX512F cvtps2ph/cvtph2ps instructions.
   */
  template <__vec_builtin _UV, __vec_builtin _TV, _ArchTraits _Traits = {}>
    [[__gnu__::__always_inline__]]
    constexpr _UV
    __vec_cast(_TV __v)
    {
      static_assert(__width_of<_UV> == __width_of<_TV>);
#if _GLIBCXX_X86
      using _Up = __vec_value_type<_UV>;
      using _Tp = __vec_value_type<_TV>;
      constexpr bool __to_f16 = is_same_v<_Up, _Float16>;
      constexpr bool __from_f16 = is_same_v<_Tp, _Float16>;
      constexpr bool __needs_f16c = _Traits._M_have_f16c() && !_Traits._M_have_avx512fp16()
                                      && (__to_f16 || __from_f16);
      if (__needs_f16c && !__is_const_known(__v))
        { // Work around PR121688
          if constexpr (__needs_f16c)
            return __x86_cvt_f16c<_UV>(__v);
        }
      if constexpr (is_floating_point_v<_Tp> && is_integral_v<_Up>
                      && sizeof(_UV) < sizeof(_TV) && sizeof(_Up) < sizeof(int))
        {
          using _Ip = __integer_from<std::min(sizeof(int), sizeof(_Tp))>;
          using _IV = __vec_builtin_type<_Ip, __width_of<_TV>>;
          return __vec_cast<_UV>(__vec_cast<_IV>(__v));
        }
#endif
      return __builtin_convertvector(__v, _UV);
    }
 
  /** @internal
   * Overload of the above cast function that determines the destination vector type from a given
   * element type @p _Up and the `__width_of` the argument type.
   *
   * Calls the above overload.
   */
  template <__vectorizable _Up, __vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr __vec_builtin_type<_Up, __width_of<_TV>>
    __vec_cast(_TV __v)
    { return __vec_cast<__vec_builtin_type<_Up, __width_of<_TV>>>(__v); }
 
  /** @internal
   * As above, but with additional precondition on possible values of the argument.
   *
   * Precondition: __k[i] is either 0 or -1 for all i.
   */
  template <__vec_builtin _UV, __vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr _UV
    __vec_mask_cast(_TV __k)
    {
      static_assert(signed_integral<__vec_value_type<_UV>>);
      static_assert(signed_integral<__vec_value_type<_TV>>);
      // TODO: __builtin_convertvector cannot be optimal because it doesn't consider input and
      // output can only be 0 or -1.
      return __builtin_convertvector(__k, _UV);
    }
 
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr _TV
    __vec_xor(_TV __a, _TV __b)
    {
      using _Tp = __vec_value_type<_TV>;
      if constexpr (is_floating_point_v<_Tp>)
        {
          using _UV = __vec_builtin_type<__integer_from<sizeof(_Tp)>, __width_of<_TV>>;
          return __builtin_bit_cast(
                   _TV, __builtin_bit_cast(_UV, __a) ^ __builtin_bit_cast(_UV, __b));
        }
      else
        return __a ^ __b;
    }
 
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr _TV
    __vec_or(_TV __a, _TV __b)
    {
      using _Tp = __vec_value_type<_TV>;
      if constexpr (is_floating_point_v<_Tp>)
        {
          using _UV = __vec_builtin_type<__integer_from<sizeof(_Tp)>, __width_of<_TV>>;
          return __builtin_bit_cast(
                   _TV, __builtin_bit_cast(_UV, __a) | __builtin_bit_cast(_UV, __b));
        }
      else
        return __a | __b;
    }
 
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr _TV
    __vec_and(_TV __a, _TV __b)
    {
      using _Tp = __vec_value_type<_TV>;
      if constexpr (is_floating_point_v<_Tp>)
        {
          using _UV = __vec_builtin_type<__integer_from<sizeof(_Tp)>, __width_of<_TV>>;
          return __builtin_bit_cast(
                   _TV, __builtin_bit_cast(_UV, __a) & __builtin_bit_cast(_UV, __b));
        }
      else
        return __a & __b;
    }
 
  /** @internal
   * Returns the bit-wise and of not @p __a and @p __b.
   *
   * Use __vec_and(__vec_not(__a), __b) unless an andnot instruction is necessary for optimization.
   *
   * @see __vec_andnot in simd_x86.h
   */
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr _TV
    __vec_andnot(_TV __a, _TV __b)
    {
      using _Tp = __vec_value_type<_TV>;
      using _UV = __vec_builtin_type<__integer_from<sizeof(_Tp)>, __width_of<_TV>>;
      return __builtin_bit_cast(
               _TV, ~__builtin_bit_cast(_UV, __a) & __builtin_bit_cast(_UV, __b));
    }
 
  template <__vec_builtin _TV>
    [[__gnu__::__always_inline__]]
    constexpr _TV
    __vec_not(_TV __a)
    {
      using _Tp = __vec_value_type<_TV>;
      using _UV = __vec_builtin_type_bytes<__integer_from<sizeof(_Tp)>, sizeof(_TV)>;
      if constexpr (is_floating_point_v<__vec_value_type<_TV>>)
        return __builtin_bit_cast(_TV, ~__builtin_bit_cast(_UV, __a));
      else
        return ~__a;
    }
 
  /**
   * An object of given type where only the sign bits are 1.
   */
  template <__vec_builtin _V>
    requires std::floating_point<__vec_value_type<_V>>
    constexpr _V _S_signmask = __vec_xor(_V() + 1, _V() - 1);
 
  template <__vec_builtin _TV, int _Np = __width_of<_TV>,
            typename = make_integer_sequence<int, _Np>>
    struct _VecOps;
 
  template <__vec_builtin _TV, int _Np, int... _Is>
    struct _VecOps<_TV, _Np, integer_sequence<int, _Is...>>
    {
      static_assert(_Np <= __width_of<_TV>);
 
      using _Tp = __vec_value_type<_TV>;
 
      using _HV = __half_vec_builtin_t<__conditional_t<_Np >= 2, _TV, __double_vec_builtin_t<_TV>>>;
 
      [[__gnu__::__always_inline__]]
      static constexpr _TV
      _S_broadcast_to_even(_Tp __init)
      { return _TV {((_Is & 1) == 0 ? __init : _Tp())...}; }
 
      [[__gnu__::__always_inline__]]
      static constexpr _TV
      _S_broadcast_to_odd(_Tp __init)
      { return _TV {((_Is & 1) == 1 ? __init : _Tp())...}; }
 
      [[__gnu__::__always_inline__]]
      static constexpr bool
      _S_all_of(_TV __k) noexcept
      { return (... && (__k[_Is] != 0)); }
 
      [[__gnu__::__always_inline__]]
      static constexpr bool
      _S_any_of(_TV __k) noexcept
      { return (... || (__k[_Is] != 0)); }
 
      [[__gnu__::__always_inline__]]
      static constexpr bool
      _S_none_of(_TV __k) noexcept
      { return (... && (__k[_Is] == 0)); }
 
      template <typename _Offset = integral_constant<int, 0>>
      [[__gnu__::__always_inline__]]
      static constexpr _TV
      _S_extract(__vec_builtin auto __x, _Offset = {})
      {
        static_assert(is_same_v<__vec_value_type<_TV>, __vec_value_type<decltype(__x)>>);
        return __builtin_shufflevector(__x, decltype(__x)(), (_Is + _Offset::value)...);
      }
 
      // swap neighboring elements
      [[__gnu__::__always_inline__]]
      static constexpr _TV
      _S_swap_neighbors(_TV __x)
      { return __builtin_shufflevector(__x, __x, (_Is ^ 1)...); }
 
      // duplicate even indexed elements, dropping the odd ones
      [[__gnu__::__always_inline__]]
      static constexpr _TV
      _S_dup_even(_TV __x)
      { return __builtin_shufflevector(__x, __x, (_Is & ~1)...); }
 
      // duplicate odd indexed elements, dropping the even ones
      [[__gnu__::__always_inline__]]
      static constexpr _TV
      _S_dup_odd(_TV __x)
      { return __builtin_shufflevector(__x, __x, (_Is | 1)...); }
 
      [[__gnu__::__always_inline__]]
      static constexpr void
      _S_overwrite_even_elements(_TV& __x, _HV __y) requires (_Np > 1)
      {
        constexpr __simd_size_type __n = __width_of<_TV>;
        __x = __builtin_shufflevector(__x,
#ifdef _GLIBCXX_CLANG
                                      __vec_concat(__y, __y),
#else
                                      __y,
#endif
                                      ((_Is & 1) == 0 ? __n + _Is / 2 : _Is)...);
      }
 
      [[__gnu__::__always_inline__]]
      static constexpr void
      _S_overwrite_even_elements(_TV& __xl, _TV& __xh, _TV __y)
      {
        constexpr __simd_size_type __nl = __width_of<_TV>;
        constexpr __simd_size_type __nh = __nl * 3 / 2;
        __xl = __builtin_shufflevector(__xl, __y, ((_Is & 1) == 0 ? __nl + _Is / 2 : _Is)...);
        __xh = __builtin_shufflevector(__xh, __y, ((_Is & 1) == 0 ? __nh + _Is / 2 : _Is)...);
      }
 
      [[__gnu__::__always_inline__]]
      static constexpr void
      _S_overwrite_odd_elements(_TV& __x, _HV __y) requires (_Np > 1)
      {
        constexpr __simd_size_type __n = __width_of<_TV>;
        __x = __builtin_shufflevector(__x,
#ifdef _GLIBCXX_CLANG
                                      __vec_concat(__y, __y),
#else
                                      __y,
#endif
                                      ((_Is & 1) == 1 ? __n + _Is / 2 : _Is)...);
      }
 
      [[__gnu__::__always_inline__]]
      static constexpr void
      _S_overwrite_odd_elements(_TV& __xl, _TV& __xh, _TV __y)
      {
        constexpr __simd_size_type __nl = __width_of<_TV>;
        constexpr __simd_size_type __nh = __nl * 3 / 2;
        __xl = __builtin_shufflevector(__xl, __y, ((_Is & 1) == 1 ? __nl + _Is / 2 : _Is)...);
        __xh = __builtin_shufflevector(__xh, __y, ((_Is & 1) == 1 ? __nh + _Is / 2 : _Is)...);
      }
 
      // true if all elements are know to be equal to __ref at compile time
      [[__gnu__::__always_inline__]]
      static constexpr bool
      _S_is_const_known_equal_to(_TV __x, _Tp __ref)
      { return (__is_const_known_equal_to(__x[_Is], __ref) && ...); }
 
    };
} // namespace simd
_GLIBCXX_END_NAMESPACE_VERSION
} // namespace std
 
#pragma GCC diagnostic pop
#endif // C++26
#endif // _GLIBCXX_VEC_OPS_H