native_f32x4_sse.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 <span>
#include <array>
#include <ostream>
 
namespace hi { inline namespace v1 {
 
#ifdef HI_HAS_SSE
 
template<>
struct native_simd<float, 4> {
    using value_type = float;
    constexpr static size_t size = 4;
 
    using array_type = std::array<value_type, size>;
    using register_type = __m128;
 
    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_ps()) {}
 
    [[nodiscard]] explicit native_simd(register_type other) noexcept : v(other) {}
 
    [[nodiscard]] explicit operator register_type() const noexcept
    {
        return v;
    }
 
    [[nodiscard]] explicit native_simd(value_type a) noexcept : v(_mm_set_ss(a)) {}
 
    [[nodiscard]] native_simd(value_type a, value_type b, value_type c = value_type{0}, value_type d = value_type{0}) noexcept :
        v(_mm_set_ps(d, c, b, a))
    {
    }
 
    [[nodiscard]] explicit native_simd(value_type const *other) noexcept : v(_mm_loadu_ps(other)) {}
 
    void store(value_type *out) const noexcept
    {
        hi_axiom_not_null(out);
        _mm_storeu_ps(out, v);
    }
 
    [[nodiscard]] explicit native_simd(void const *other) noexcept : v(_mm_loadu_ps(static_cast<value_type const *>(other))) {}
 
    void store(void *out) const noexcept
    {
        hi_axiom_not_null(out);
        _mm_storeu_ps(static_cast<value_type *>(out), v);
    }
 
    [[nodiscard]] explicit native_simd(std::span<value_type const> other) noexcept
    {
        hi_axiom(other.size() >= size);
        v = _mm_loadu_ps(other.data());
    }
 
    void store(std::span<value_type> out) const noexcept
    {
        hi_axiom(out.size() >= size);
        _mm_storeu_ps(out.data(), v);
    }
 
    [[nodiscard]] explicit native_simd(array_type other) noexcept : v(_mm_loadu_ps(other.data())) {}
 
    [[nodiscard]] explicit operator array_type() const noexcept
    {
        auto r = array_type{};
        _mm_storeu_ps(r.data(), v);
        return r;
    }
 
#ifdef HI_HAS_SSE2
    [[nodiscard]] explicit native_simd(native_simd<int32_t, 4> const& a) noexcept;
#endif
#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_ps(a)};
    }
 
    [[nodiscard]] static native_simd broadcast(native_simd a) noexcept
    {
#ifdef HI_HAS_AVX2
        return native_simd{_mm_broadcastss_ps(a.v)};
#else
        return native_simd{_mm_shuffle_ps(a.v, a.v, 0b00'00'00'00)};
#endif
    }
 
    [[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{_mm_castsi128_ps(tmp)};
    }
 
    [[nodiscard]] static native_simd ones() noexcept
    {
#ifdef HI_HAS_SSE2
        auto ones = _mm_undefined_si128();
        ones = _mm_cmpeq_epi32(ones, ones);
        return native_simd{_mm_castsi128_ps(ones)};
#else
        auto ones = _mm_setzero_ps();
        ones = _mm_cmpeq_ps(ones, ones);
        return native_simd{ones};
#endif
    }
 
    [[nodiscard]] size_t mask() const noexcept
    {
        return narrow_cast<size_t>(_mm_movemask_ps(v));
    }
 
    [[nodiscard]] friend bool equal(native_simd a, native_simd b) noexcept
    {
#ifdef HI_HAS_SSE2
        return _mm_movemask_epi8(_mm_cmpeq_epi32(_mm_castps_si128(a.v), _mm_castps_si128(b.v))) == 0b1111'1111'1111'1111;
#else
        return static_cast<array_type>(a) == static_cast<array_type>(b);
#endif
    }
 
    [[nodiscard]] friend native_simd
    almost_eq(native_simd a, native_simd b, value_type epsilon = std::numeric_limits<value_type>::epsilon()) noexcept
    {
        hilet abs_diff = abs(a - b);
        return abs_diff < broadcast(epsilon);
    }
 
    [[nodiscard]] friend bool
    almost_equal(native_simd a, native_simd b, value_type epsilon = std::numeric_limits<value_type>::epsilon())
    {
        return almost_eq(a, b, epsilon).mask() == 0b1111;
    }
 
    [[nodiscard]] friend native_simd operator==(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmpeq_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator!=(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmpneq_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator<(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmplt_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator>(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmpgt_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator<=(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmple_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator>=(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_cmpge_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator+(native_simd a) noexcept
    {
        return a;
    }
 
    [[nodiscard]] friend native_simd operator+(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_add_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator-(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_sub_ps(a.v, b.v)};
    }
 
    [[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_mul_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator/(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_div_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator&(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_and_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator|(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_or_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator^(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_xor_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd operator~(native_simd a) noexcept
    {
        return not_and(a, ones());
    }
 
    [[nodiscard]] friend native_simd min(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_min_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd max(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_max_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend native_simd abs(native_simd a) noexcept
    {
        return not_and(broadcast(-0.0f), a);
    }
 
#ifdef HI_HAS_SSE4_1
    [[nodiscard]] friend native_simd floor(native_simd a) noexcept
    {
        return native_simd{_mm_floor_ps(a.v)};
    }
#endif
 
#ifdef HI_HAS_SSE4_1
    [[nodiscard]] friend native_simd ceil(native_simd a) noexcept
    {
        return native_simd{_mm_ceil_ps(a.v)};
    }
#endif
 
#ifdef HI_HAS_SSE4_1
    template<native_rounding_mode Rounding = native_rounding_mode::current>
    [[nodiscard]] friend native_simd round(native_simd a) noexcept
    {
        return native_simd{_mm_round_ps(a.v, to_underlying(Rounding))};
    }
#endif
 
    [[nodiscard]] friend native_simd rcp(native_simd a) noexcept
    {
        return native_simd{_mm_rcp_ps(a.v)};
    }
 
    [[nodiscard]] friend native_simd sqrt(native_simd a) noexcept
    {
        return native_simd{_mm_sqrt_ps(a.v)};
    }
 
    [[nodiscard]] friend native_simd rsqrt(native_simd a) noexcept
    {
        return native_simd{_mm_rsqrt_ps(a.v)};
    }
 
    template<size_t Mask>
    [[nodiscard]] friend native_simd set_zero(native_simd a) noexcept
    {
        static_assert(Mask <= 0b1111);
        if constexpr (Mask == 0b0000) {
            return a;
        } else if constexpr (Mask == 0b1111) {
            return {};
        } else {
#ifdef HI_HAS_SSE4_1
            return native_simd{_mm_insert_ps(a.v, 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_ps(a.v, _mm_set1_ps(b), narrow_cast<int>(Index << 4))};
#else
        hilet mask = from_mask(1_uz << Index);
        return not_and(mask, a) | (mask & broadcast(b));
#endif
    }
 
    template<size_t SrcIndex, size_t DstIndex>
    [[nodiscard]] friend native_simd insert(native_simd a, native_simd b) noexcept
    {
        static_assert(SrcIndex < size);
        static_assert(DstIndex < size);
#ifdef HI_HAS_SSE4_1
        return native_simd{_mm_insert_ps(a.v, b.v, (SrcIndex << 6) | (DstIndex << 4))};
#else
        return insert<DstIndex>(a, get<SrcIndex>(b));
#endif
    }
 
    template<size_t Index>
    [[nodiscard]] friend value_type get(native_simd a) noexcept
    {
        static_assert(Index < size);
 
        hilet tmp = _mm_shuffle_ps(a.v, a.v, Index);
        return _mm_cvtss_f32(tmp);
    }
 
    template<size_t Mask>
    [[nodiscard]] friend native_simd blend(native_simd a, native_simd b) noexcept
    {
        static_assert(Mask <= 0b1111);
 
        if constexpr (Mask == 0b0000) {
            return a;
        } else if constexpr (Mask == 0b1111) {
            return b;
        } else {
#ifdef HI_HAS_SSE4_1
            return native_simd{_mm_blend_ps(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
    {
        static_assert(SourceElements.size() == size);
        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 {
#ifdef HI_HAS_AVX
            return native_simd{_mm_permute_ps(a.v, order)};
#else
            return native_simd{_mm_shuffle_ps(a.v, a.v, order)};
#endif
        }
    }
 
    [[nodiscard]] friend native_simd permute(native_simd a, native_simd<int32_t, 4> const& source_elements) noexcept;
 
    template<fixed_string SourceElements>
    [[nodiscard]] friend native_simd swizzle(native_simd a) noexcept
    {
        static_assert(SourceElements.size() == size);
        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_ps(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_ps(a.v, b.v)};
    }
#endif
 
    [[nodiscard]] friend native_simd horizontal_sum(native_simd a) noexcept
    {
        hilet 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);
#ifdef HI_HAS_SSE4_1
        return native_simd{_mm_dp_ps(a.v, b.v, (SourceMask << 4) | 0b1111)};
#else
        return horizontal_sum(set_zero<~SourceMask & 0b1111>(a * b));
#endif
    }
 
#ifdef HI_HAS_SSE3
    [[nodiscard]] friend native_simd interleaved_sub_add(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_addsub_ps(a.v, b.v)};
    }
#endif
 
    [[nodiscard]] friend native_simd not_and(native_simd a, native_simd b) noexcept
    {
        return native_simd{_mm_andnot_ps(a.v, b.v)};
    }
 
    [[nodiscard]] friend std::array<native_simd, 4> transpose(native_simd a, native_simd b, native_simd c, native_simd d) noexcept
    {
        _MM_TRANSPOSE4_PS(a.v, b.v, c.v, d.v);
        return {a, b, c, d};
    }
 
    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 {};
 
        } else if constexpr ((one_mask | alpha_mask) == 0b1111) {
            return broadcast(1.0f);
 
        } else {
            return native_simd{
                to_bool(one_mask & 0b0001) ? 1.0f : 0.0f,
                to_bool(one_mask & 0b0010) ? 1.0f : 0.0f,
                to_bool(one_mask & 0b0100) ? 1.0f : 0.0f,
                to_bool(one_mask & 0b1000) ? 1.0f : 0.0f};
        }
    }
};
 
#endif
}} // namespace hi::v1