xsimd/arch/xsimd_avx2_128.hpp

0001 /***************************************************************************
0002  * Copyright (c) Johan Mabille, Sylvain Corlay, Wolf Vollprecht and         *
0003  * Martin Renou                                                             *
0004  * Copyright (c) QuantStack                                                 *
0005  * Copyright (c) Serge Guelton                                              *
0006  * Copyright (c) Marco Barbone                                              *
0007  *                                                                          *
0008  * Distributed under the terms of the BSD 3-Clause License.                 *
0009  *                                                                          *
0010  * The full license is in the file LICENSE, distributed with this software. *
0011  ****************************************************************************/
0012
0013 #ifndef XSIMD_AVX2_128_HPP
0014 #define XSIMD_AVX2_128_HPP
0015
0016 #include "../types/xsimd_avx2_register.hpp"
0017 #include "../types/xsimd_batch_constant.hpp"
0018
0019 #include <type_traits>
0020
0021 namespace xsimd
0022 {
0023     namespace kernel
0024     {
0025         using namespace types;
0026
0027         // Defined in xsimd_avx512vl_128.hpp (included after this header). The
0028         // masked load/store half-fold below forwards to the 128-bit sized-batch
0029         // arch, which is avx512vl_128 when the 256-bit one stays in the AVX2
0030         // lineage (e.g. avxvnni). That unqualified dependent call resolves by
0031         // ordinary lookup here, so the avx512vl_128 overloads must be visible at
0032         // this point; declarations only, never instantiated unless AVX512VL is on.
0033         template <class A, class T, bool... V, class Mode,
0034                   typename = std::enable_if_t<std::is_arithmetic<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
0035         XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, V...> mask, convert<T>, Mode, requires_arch<avx512vl_128>) noexcept;
0036         template <class A, class T, class Mode,
0037                   typename = std::enable_if_t<std::is_arithmetic<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
0038         XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<avx512vl_128>) noexcept;
0039         template <class A, class T, bool... V, class Mode,
0040                   typename = std::enable_if_t<std::is_arithmetic<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
0041         XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, V...> mask, Mode, requires_arch<avx512vl_128>) noexcept;
0042         template <class A, class T, class Mode,
0043                   typename = std::enable_if_t<std::is_arithmetic<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
0044         XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<avx512vl_128>) noexcept;
0045
0046         // select
0047         template <class A, class T, bool... Values, class = std::enable_if_t<std::is_integral<T>::value>>
0048         XSIMD_INLINE batch<T, A> select(batch_bool_constant<T, A, Values...> const&, batch<T, A> const& true_br, batch<T, A> const& false_br, requires_arch<avx2_128>) noexcept
0049         {
0050             constexpr int mask = batch_bool_constant<T, A, Values...>::mask();
0051             XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
0052             {
0053                 return _mm_blend_epi32(false_br, true_br, mask);
0054             }
0055             else
0056             {
0057                 return select(batch_bool_constant<T, A, Values...>(), true_br, false_br, avx_128 {});
0058             }
0059         }
0060
0061         // bitwise_lshift
0062         template <class A, class T, class = std::enable_if_t<std::is_integral<T>::value>>
0063         XSIMD_INLINE batch<T, A> bitwise_lshift(batch<T, A> const& self, batch<T, A> const& other, requires_arch<avx2_128>) noexcept
0064         {
0065             XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
0066             {
0067                 return _mm_sllv_epi32(self, other);
0068             }
0069             else XSIMD_IF_CONSTEXPR(sizeof(T) == 8)
0070             {
0071                 return _mm_sllv_epi64(self, other);
0072             }
0073             else
0074             {
0075                 return bitwise_lshift(self, other, avx {});
0076             }
0077         }
0078
0079         // bitwise_rshift
0080         template <class A, class T, class = std::enable_if_t<std::is_integral<T>::value>>
0081         XSIMD_INLINE batch<T, A> bitwise_rshift(batch<T, A> const& self, batch<T, A> const& other, requires_arch<avx2_128>) noexcept
0082         {
0083             if (std::is_signed<T>::value)
0084             {
0085                 XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
0086                 {
0087                     return _mm_srav_epi32(self, other);
0088                 }
0089                 else
0090                 {
0091                     return bitwise_rshift(self, other, avx_128 {});
0092                 }
0093             }
0094             else
0095             {
0096                 XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
0097                 {
0098                     return _mm_srlv_epi32(self, other);
0099                 }
0100                 else XSIMD_IF_CONSTEXPR(sizeof(T) == 8)
0101                 {
0102                     return _mm_srlv_epi64(self, other);
0103                 }
0104                 else
0105                 {
0106                     return bitwise_rshift(self, other, avx_128 {});
0107                 }
0108             }
0109         }
0110
0111         // load_masked / store_masked: native 128-bit integer masked memory.
0112         // Tagged on avx2_128 because vpmaskmov* needs AVX2; an AVX1-only build
0113         // routes integer masked memory through the float path in
0114         // xsimd_common_memory.hpp. 8/16-bit fall back to the common scalar path.
0115         namespace detail
0116         {
0117             template <class T>
0118             XSIMD_INLINE __m128i maskload_avx2_128(T const* mem, __m128i mask) noexcept
0119             {
0120                 XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
0121                 {
0122                     return _mm_maskload_epi32(reinterpret_cast<int const*>(mem), mask);
0123                 }
0124                 else
0125                 {
0126                     return _mm_maskload_epi64(reinterpret_cast<long long const*>(mem), mask);
0127                 }
0128             }
0129
0130             template <class T>
0131             XSIMD_INLINE void maskstore_avx2_128(T* mem, __m128i mask, __m128i src) noexcept
0132             {
0133                 XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
0134                 {
0135                     _mm_maskstore_epi32(reinterpret_cast<int*>(mem), mask, src);
0136                 }
0137                 else
0138                 {
0139                     _mm_maskstore_epi64(reinterpret_cast<long long*>(mem), mask, src);
0140                 }
0141             }
0142         }
0143
0144         // Masks that lower to plain moves go to the sse2 float/double kernels
0145         // (int bitcast to same-width float); the rest keep the native vpmaskmov.
0146         template <class A, class T, bool... Values, class Mode,
0147                   typename = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
0148         XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, Values...> mask, convert<T>, Mode, requires_arch<avx2_128>) noexcept
0149         {
0150             XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
0151             {
0152                 return detail::plain_move_load<sse2>(mem, mask, convert<T> {}, Mode {});
0153             }
0154             else
0155             {
0156                 return load_masked(mem, mask.as_batch_bool(), convert<T> {}, Mode {}, avx2_128 {});
0157             }
0158         }
0159
0160         template <class A, class T, bool... Values, class Mode,
0161                   typename = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
0162         XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, Values...> mask, Mode, requires_arch<avx2_128>) noexcept
0163         {
0164             XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
0165             {
0166                 detail::plain_move_store<sse2>(mem, src, mask, Mode {});
0167             }
0168             else
0169             {
0170                 store_masked(mem, src, mask.as_batch_bool(), Mode {}, avx2_128 {});
0171             }
0172         }
0173
0174         template <class A, class T, class Mode>
0175         XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), batch<T, A>>
0176         load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<avx2_128>) noexcept
0177         {
0178             return detail::maskload_avx2_128(mem, __m128i(mask));
0179         }
0180
0181         template <class A, class T, class Mode>
0182         XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), void>
0183         store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<avx2_128>) noexcept
0184         {
0185             detail::maskstore_avx2_128(mem, __m128i(mask), __m128i(src));
0186         }
0187
0188         // gather
0189         template <class T, class A, class U, detail::enable_sized_integral_t<T, 4> = 0, detail::enable_sized_integral_t<U, 4> = 0>
0190         XSIMD_INLINE batch<T, A> gather(batch<T, A> const&, T const* src, batch<U, A> const& index,
0191                                         kernel::requires_arch<avx2_128>) noexcept
0192         {
0193             return _mm_i32gather_epi32(reinterpret_cast<const int*>(src), index, sizeof(T));
0194         }
0195
0196         template <class T, class A, class U, detail::enable_sized_integral_t<T, 8> = 0, detail::enable_sized_integral_t<U, 8> = 0>
0197         XSIMD_INLINE batch<T, A> gather(batch<T, A> const&, T const* src, batch<U, A> const& index,
0198                                         kernel::requires_arch<avx2_128>) noexcept
0199         {
0200             return _mm_i64gather_epi64(reinterpret_cast<const long long int*>(src), index, sizeof(T));
0201         }
0202
0203         template <class A, class U,
0204                   detail::enable_sized_integral_t<U, 4> = 0>
0205         XSIMD_INLINE batch<float, A> gather(batch<float, A> const&, float const* src,
0206                                             batch<U, A> const& index,
0207                                             kernel::requires_arch<avx2_128>) noexcept
0208         {
0209             return _mm_i32gather_ps(src, index, sizeof(float));
0210         }
0211
0212         template <class A, class U, detail::enable_sized_integral_t<U, 8> = 0>
0213         XSIMD_INLINE batch<double, A> gather(batch<double, A> const&, double const* src,
0214                                              batch<U, A> const& index,
0215                                              requires_arch<avx2_128>) noexcept
0216         {
0217             return _mm_i64gather_pd(src, index, sizeof(double));
0218         }
0219     }
0220 }
0221
0222 #endif