native_i32x4_sse2.hpp

// Copyright Take Vos 2022, 2023.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt)
 
#pragma once
 
#include "native_simd_utility.hpp"
#include "../utility/module.hpp"
#include <array>
#include <ostream>
 
hi_warning_push();
// Ignore "C26490: Don't use reinterpret_cast", needed for intrinsic loads and stores.
hi_warning_ignore_msvc(26490);
 
namespace hi { inline namespace v1 {
 
#ifdef HI_HAS_SSE2
 
template<>
struct native_simd<int32_t, 4> {
    using value_type = int32_t;
    constexpr static size_t size = 4;
    using register_type = __m128i;
    using array_type = std::array<value_type, size>;
 
    register_type v;
 
    native_simd(native_simd const&) noexcept = default;
    native_simd(native_simd&&) noexcept = default;
    native_simd& operator=(native_simd const&) noexcept = default;
    native_simd& operator=(native_simd&&) noexcept = default;
 
    native_simd() noexcept : v(_mm_setzero_si128()) {}
 
    [[nodiscard]] explicit native_simd(register_type other) noexcept : v(other) {}
 
    [[nodiscard]] explicit operator register_type() const noexcept
    {
        return v;
    }
 
    [[nodiscard]] native_simd(
        value_type a,
        value_type b = value_type{0},
        value_type c = value_type{0},
        value_type d = value_type{0}) noexcept :
        v(_mm_set_epi32(d, c, b, a))
    {
    }
 
    [[nodiscard]] explicit native_simd(value_type const *other) noexcept :
        v(_mm_loadu_si128(reinterpret_cast<register_type const *>(other)))
    {
    }
 
    void store(value_type *out) const noexcept
    {
        hi_axiom_not_null(out);
        _mm_storeu_si128(reinterpret_cast<register_type *>(out), v);
    }
 
    [[nodiscard]] explicit native_simd(void const *other) noexcept : v(_mm_loadu_si128(static_cast<register_type const *>(other)))
    {
    }
 
    void store(void *out) const noexcept
    {
        hi_axiom_not_null(out);
        _mm_storeu_si128(static_cast<register_type *>(out), v);
    }
 
    [[nodiscard]] explicit native_simd(std::span<value_type const> other) noexcept
    {
        hi_axiom(other.size() >= size);
        v = _mm_loadu_si128(reinterpret_cast<register_type const *>(other.data()));
    }
 
    void store(std::span<value_type> out) const noexcept
    {
        hi_axiom(out.size() >= size);
        _mm_storeu_si128(reinterpret_cast<register_type *>(out.data()), v);
    }
 
    [[nodiscard]] explicit native_simd(array_type other) noexcept :
        v(_mm_loadu_si128(reinterpret_cast<register_type const *>(other.data())))
    {
    }
 
    [[nodiscard]] explicit operator array_type() const noexcept
    {
        auto r = array_type{};
        _mm_storeu_si128(reinterpret_cast<register_type *>(r.data()), v);
        return r;
    }
 
    [[nodiscard]] explicit native_simd(native_simd<float, 4> const& a) noexcept;
    [[nodiscard]] explicit native_simd(native_simd<uint32_t, 4> const& a) noexcept;
#ifdef HI_HAS_AVX
    [[nodiscard]] explicit native_simd(native_simd<double, 4> const& a) noexcept;
#endif
 
    [[nodiscard]] static native_simd broadcast(value_type a) noexcept
    {
        return native_simd{_mm_set1_epi32(a)};
    }
 
    [[nodiscard]] static native_simd broadcast(native_simd a) noexcept
    {
#ifdef HI_HAS_AVX2
        return native_simd{_mm_broadcastd_epi32(a.v)};
#else
        return native_simd{_mm_shuffle_epi32(a.v, 0b00'00'00'00)};
#endif
    }
 
    [[nodiscard]] static native_simd ones() noexcept
    {
        hilet tmp = _mm_undefined_si128();
        return native_simd{_mm_cmpeq_epi32(tmp, tmp)};
    }
 
    template<size_t Mask>
    [[nodiscard]] static native_simd from_mask() noexcept
    {
        return native_simd{
            to_bool(Mask & 0b0001) ? static_cast<value_type>(0xffff'ffff) : 0,
            to_bool(Mask & 0b0010) ? static_cast<value_type>(0xffff'ffff) : 0,
            to_bool(Mask & 0b0100) ? static_cast<value_type>(0xffff'ffff) : 0,
            to_bool(Mask & 0b1000) ? static_cast<value_type>(0xffff'ffff) : 0};
    }
 
    [[nodiscard]] static native_simd from_mask(size_t a) noexcept
    {
        hi_axiom(a <= 0b1111);
 
        uint64_t a_ = a;
 
        a_ <<= 31;
        auto tmp = _mm_cvtsi32_si128(truncate<uint32_t>(a_));
        a_ >>= 1;
        tmp = _mm_insert_epi32(tmp, truncate<uint32_t>(a_), 1);
        a_ >>= 1;
        tmp = _mm_insert_epi32(tmp, truncate<uint32_t>(a_), 2);
        a_ >>= 1;
        tmp = _mm_insert_epi32(tmp, truncate<uint32_t>(a_), 3);
 
        tmp = _mm_srai_epi32(tmp, 31);
        return native_simd{tmp};
    }
 
    [[nodiscard]] size_t mask() const noexcept
    {
        return narrow_cast<size_t>(_mm_movemask_ps(_mm_castsi128_ps(v)));
    }
 
    [[nodiscard]] friend bool equal(native_simd a, native_simd b) noexcept
    {
        return (a == b).mask() == 0b1111;
    }
 
    [[nodiscard]] friend native_simd operator==(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmpeq_epi32(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator!=(native_simd a, native_simd b) noexcept
    {
        return ~(a == b);
    }
 
    [[nodiscard]] friend native_simd operator<(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmplt_epi32(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator>(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmpgt_epi32(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator<=(native_simd a, native_simd b) noexcept
    {
        return ~(a > b);
    }
 
    [[nodiscard]] friend native_simd operator>=(native_simd a, native_simd b) noexcept
    {
        return ~(a < b);
    }
 
    [[nodiscard]] friend native_simd operator+(native_simd a) noexcept
    {
        return a;
    }
 
    [[nodiscard]] friend native_simd operator-(native_simd a) noexcept
    {
        return native_simd{} - a;
    }
 
    [[nodiscard]] friend native_simd operator+(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_add_epi32(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator-(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_sub_epi32(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator*(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_mullo_epi32(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator&(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_and_si128(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator|(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_or_si128(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator^(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_xor_si128(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator~(native_simd a) noexcept
    {
        auto ones = _mm_undefined_si128();
        ones = _mm_cmpeq_epi32(ones, ones);
        return native_simd{_mm_andnot_si128(a.v, ones)};
    }
 
    [[nodiscard]] friend native_simd operator<<(native_simd a, unsigned int b) noexcept
    {
        hi_axiom_bounds(b, sizeof(value_type) * CHAR_BIT);
        return native_simd{_mm_slli_epi32(a.v, b)};
    }
 
    [[nodiscard]] friend native_simd operator>>(native_simd a, unsigned int b) noexcept
    {
        hi_axiom_bounds(b, sizeof(value_type) * CHAR_BIT);
        return native_simd{_mm_srai_epi32(a.v, b)};
    }
 
    [[nodiscard]] friend native_simd min(native_simd a, native_simd b) noexcept
    {
#if HI_HAS_SSE4_1
        return native_simd{_mm_min_epi32(a.v, b.v)};
#else
        hilet mask = a < b;
        return (mask & a) | not_and(mask, b);
#endif
    }
 
    [[nodiscard]] friend native_simd max(native_simd a, native_simd b) noexcept
    {
#if HI_HAS_SSE4_1
        return native_simd{_mm_max_epi32(a.v, b.v)};
#else
        hilet mask = a > b;
        return (mask & a) | not_and(mask, b);
#endif
    }
 
    [[nodiscard]] friend native_simd abs(native_simd a) noexcept
    {
#if HI_HAS_SSSE3
        return native_simd{_mm_abs_epi32(a.v)};
#else
        hilet mask = a >= native_simd{};
        return (mask & a) | not_and(mask, -a);
#endif
    }
 
    template<size_t Mask>
    [[nodiscard]] friend native_simd set_zero(native_simd a) noexcept
    {
        static_assert(Mask <= 0b1111);
#ifdef HI_HAS_SSE4_1
        return native_simd{_mm_castps_si128(_mm_insert_ps(_mm_castsi128_ps(a.v), _mm_castsi128_ps(a.v), Mask))};
#else
        hilet mask = from_mask<Mask>();
        return not_and(mask, a);
#endif
    }
 
    template<size_t Index>
    [[nodiscard]] friend native_simd insert(native_simd a, value_type b) noexcept
    {
        static_assert(Index < 4);
 
#ifdef HI_HAS_SSE4_1
        return native_simd{_mm_insert_epi32(a.v, b, Index)};
#else
        hilet mask = from_mask<1_uz << Index>();
        return not_and(mask, a) | (mask & broadcast(b));
#endif
    }
 
    template<size_t Index>
    [[nodiscard]] friend value_type get(native_simd a) noexcept
    {
#ifdef HI_HAS_SSE4_1
        return _mm_extract_epi32(a.v, Index);
#else
        auto r = static_cast<array_type>(a);
        return std::get<Index>(r);
#endif
    }
 
    template<size_t Mask>
    [[nodiscard]] friend native_simd blend(native_simd a, native_simd b) noexcept
    {
#ifdef HI_HAS_SSE4_1
        return native_simd{_mm_blend_epi32(a.v, b.v, Mask)};
#else
        hilet mask = from_mask<Mask>();
        return not_and(mask, a) | (mask & b);
#endif
    }
 
    template<fixed_string SourceElements>
    [[nodiscard]] friend native_simd permute(native_simd a) noexcept
    {
        constexpr auto order = detail::native_swizzle_to_packed_indices<SourceElements, size>();
 
        if constexpr (order == 0b11'10'01'00) {
            return a;
        } else if constexpr (order == 0b00'00'00'00) {
            return broadcast(a);
        } else {
            return native_simd{_mm_shuffle_epi32(a.v, order)};
        }
    }
 
    template<fixed_string SourceElements>
    [[nodiscard]] friend native_simd swizzle(native_simd a) noexcept
    {
        constexpr auto one_mask = detail::native_swizzle_to_mask<SourceElements, size, '1'>();
        constexpr auto zero_mask = detail::native_swizzle_to_mask<SourceElements, size, '0'>();
        constexpr auto number_mask = one_mask | zero_mask;
 
        if constexpr (number_mask == 0b1111) {
            // Swizzle was /[01][01][01][01]/.
            return swizzle_numbers<SourceElements>();
 
        } else if constexpr (number_mask == 0b0000) {
            // Swizzle was /[^01][^01][^01][^01]/.
            return permute<SourceElements>(a);
 
#ifdef HI_HAS_SSE4_1
        } else if constexpr (number_mask == zero_mask) {
            // Swizzle was /[^1][^1][^1][^1]/.
            hilet ordered = permute<SourceElements>(a);
            return set_zero<zero_mask>(ordered);
#endif
 
        } else {
            hilet ordered = permute<SourceElements>(a);
            hilet numbers = swizzle_numbers<SourceElements>();
            return blend<number_mask>(ordered, numbers);
        }
    }
 
#ifdef HI_HAS_SSE3
    [[nodiscard]] friend native_simd horizontal_add(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_hadd_epi32(a.v, b.v)};
    }
#endif
 
#ifdef HI_HAS_SSE3
    [[nodiscard]] friend native_simd horizontal_sub(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_hsub_epi32(a.v, b.v)};
    }
#endif
 
    [[nodiscard]] friend native_simd horizontal_sum(native_simd a) noexcept
    {
        auto tmp = a + permute<"cdab">(a);
        return tmp + permute<"badc">(tmp);
    }
 
    template<size_t SourceMask>
    [[nodiscard]] friend native_simd dot_product(native_simd a, native_simd b) noexcept
    {
        static_assert(SourceMask <= 0b1111);
        return horizontal_sum(set_zero<~SourceMask & 0b1111>(a * b));
    }
 
    [[nodiscard]] friend native_simd not_and(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_andnot_si128(a.v, b.v)};
    }
 
    friend std::ostream& operator<<(std::ostream& a, native_simd b) noexcept
    {
        return a << "(" << get<0>(b) << ", " << get<1>(b) << ", " << get<2>(b) << ", " << get<3>(b) << ")";
    }
 
    template<fixed_string SourceElements>
    [[nodiscard]] static native_simd swizzle_numbers() noexcept
    {
        constexpr auto one_mask = detail::native_swizzle_to_mask<SourceElements, size, '1'>();
        constexpr auto zero_mask = detail::native_swizzle_to_mask<SourceElements, size, '0'>();
        constexpr auto number_mask = one_mask | zero_mask;
        constexpr auto alpha_mask = ~number_mask & 0b1111;
 
        if constexpr ((zero_mask | alpha_mask) == 0b1111) {
            return native_simd{_mm_setzero_si128()};
 
        } else if constexpr ((one_mask | alpha_mask) == 0b1111) {
            return native_simd{_mm_set1_epi32(1)};
 
        } else {
            return native_simd{
                to_bool(one_mask & 0b0001) ? 1 : 0,
                to_bool(one_mask & 0b0010) ? 1 : 0,
                to_bool(one_mask & 0b0100) ? 1 : 0,
                to_bool(one_mask & 0b1000) ? 1 : 0};
        }
    }
};
 
#endif
 
}} // namespace hi::v1
 
hi_warning_pop();