cp-algo/util/simd.hpp

• View this file on GitHub
• Last update: 2025-05-31 15:35:36+02:00
• Include: #include "cp-algo/util/simd.hpp"
• Link: https://stackoverflow.com/a/77376595

Required by

• cp-algo/linalg/frobenius.hpp
• cp-algo/linalg/matrix.hpp
• cp-algo/linalg/vector.hpp
• cp-algo/math/cvector.hpp
• cp-algo/math/factorials.hpp
• cp-algo/math/fft.hpp
• cp-algo/math/fft64.hpp
• cp-algo/math/multivar.hpp
• cp-algo/math/poly.hpp
• cp-algo/math/poly/impl/div.hpp
• cp-algo/math/poly/impl/euclid.hpp
• cp-algo/structures/bit_array.hpp
• cp-algo/structures/bitpack.hpp
• cp-algo/structures/fenwick_set.hpp
• cp-algo/util/bit.hpp
• cp-algo/util/compress_coords.hpp
• cp-algo/util/sort.hpp

Verified with

• Many Factorials (verify/combi/many_facts.test.cpp)
• Many Factorials (dynamic Mod) (verify/combi/many_facts_dynamic.test.cpp)
• Adjugate Matrix (verify/linalg/adj.test.cpp)
• Characteristic Polynomial (verify/linalg/characteristic.test.cpp)
• Matrix Determinant (verify/linalg/det.test.cpp)
• Counting Eulerian Circuits (verify/linalg/euler_circs.test.cpp)
• Inverse Matrix (verify/linalg/inv.test.cpp)
• Matching on General Graph (verify/linalg/matching.test.cpp)
• Pow of Matrix (verify/linalg/pow.test.cpp)
• Pow of Matrix (Frobenius) (verify/linalg/pow_fast.test.cpp)
• Matrix Product (verify/linalg/prod.test.cpp)
• Matrix Product (dynamic modint) (verify/linalg/prod_dynamic_modint.test.cpp)
• Rank of Matrix (verify/linalg/rank.test.cpp)
• Counting Spanning Trees (Directed) (verify/linalg/spanning_directed.test.cpp)
• Counting Spanning Trees (Undirected) (verify/linalg/spanning_undirected.test.cpp)
• System of Linear Equations (verify/linalg/system.test.cpp)
• Matching on General Graph (Tutte matrix) (verify/linalg/tutte.test.cpp)
• Bell Number (verify/poly/bell.test.cpp)
• Convolution mod $10^9+7$ (verify/poly/convolution107.test.cpp)
• Convolution (Mod $2^{64}$) (verify/poly/convolution64.test.cpp)
• Convolution (Large) (verify/poly/convolution_large.test.cpp)
• Convolution on the Multiplicative Monoid of $\mathbb Z/p\mathbb{Z}$ (verify/poly/convolution_mul.test.cpp)
• Division of Polynomials (verify/poly/div.test.cpp)
• Exp of Power Series (verify/poly/exp.test.cpp)
• Find Linear Recurrence (verify/poly/find_linrec.test.cpp)
• Polynomial Interpolation (verify/poly/inter.test.cpp)
• Polynomial Interpolation (Geometric Sequence) (verify/poly/inter_geo.test.cpp)
• Inv of Power Series (verify/poly/inv.test.cpp)
• Inv of Polynomials (verify/poly/inv_mod.test.cpp)
• Consecutive Terms of Linear Recursion (verify/poly/linrec_consecutive.test.cpp)
• Kth term of Linearly Recurrent Sequence (verify/poly/linrec_single.test.cpp)
• Log of Power Series (verify/poly/log.test.cpp)
• Multipoint Evaluation (verify/poly/multieval.test.cpp)
• Multipoint Evaluation (Geometric Sequence) (verify/poly/multieval_geo.test.cpp)
• Multidimensional Convolution (Truncated) (verify/poly/multivar.test.cpp)
• Conversion to Newton Basis (verify/poly/newton.test.cpp)
• Pow of Power Series (verify/poly/pow.test.cpp)
• Product of Polynomial Sequence (verify/poly/prod_sequence.test.cpp)
• Polynomial Root Finding (verify/poly/roots.test.cpp)
• Shift of Sampling Points of Polynomial (verify/poly/sampling.test.cpp)
• Sqrt of Power Series (verify/poly/sqrt.test.cpp)
• $\sum\limits_{i=0}^\infty r^i i^d$ (verify/poly/sum_exp_poly.test.cpp)
• Polynomial Taylor Shift (verify/poly/taylor.test.cpp)
• Wildcard Pattern Matching (verify/poly/wildcard.test.cpp)
• Determinant of Matrix (Mod 2) (verify/structures/bitpack/det_mod_2.test.cpp)
• Inverse Matrix (Mod 2) (verify/structures/bitpack/inv_mod_2.test.cpp)
• Matrix Product (Mod 2) (verify/structures/bitpack/prod_mod_2.test.cpp)
• System of Linear Equations (Mod 2) (verify/structures/bitpack/system_mod_2.test.cpp)
• Longest Increasing Subsequence (verify/structures/fenwick/lis.test.cpp)
• Ordered Set (verify/structures/fenwick/ordered_set.test.cpp)

Code

#ifndef CP_ALGO_UTIL_SIMD_HPP
#define CP_ALGO_UTIL_SIMD_HPP
#include <experimental/simd>
#include <cstdint>
#include <cstddef>
#include <memory>
namespace cp_algo {
    template<typename T, size_t len>
    using simd [[gnu::vector_size(len * sizeof(T))]] = T;
    using i64x4 = simd<int64_t, 4>;
    using u64x4 = simd<uint64_t, 4>;
    using u32x8 = simd<uint32_t, 8>;
    using i32x4 = simd<int32_t, 4>;
    using u32x4 = simd<uint32_t, 4>;
    using i16x4 = simd<int16_t, 4>;
    using u8x32 = simd<uint8_t, 32>;
    using dx4 = simd<double, 4>;

    [[gnu::target("avx2")]] inline dx4 abs(dx4 a) {
    return a < 0 ? -a : a;
    }

    // https://stackoverflow.com/a/77376595
    // works for ints in (-2^51, 2^51)
    static constexpr dx4 magic = dx4() + (3ULL << 51);
    [[gnu::target("avx2")]] inline i64x4 lround(dx4 x) {
        return i64x4(x + magic) - i64x4(magic);
    }
    [[gnu::target("avx2")]] inline dx4 to_double(i64x4 x) {
        return dx4(x + i64x4(magic)) - magic;
    }

    [[gnu::target("avx2")]] inline dx4 round(dx4 a) {
        return dx4{
            std::nearbyint(a[0]),
            std::nearbyint(a[1]),
            std::nearbyint(a[2]),
            std::nearbyint(a[3])
        };
    }

    [[gnu::target("avx2")]] inline u64x4 low32(u64x4 x) {
        return x & uint32_t(-1);
    }
    [[gnu::target("avx2")]] inline auto swap_bytes(auto x) {
        return decltype(x)(__builtin_shufflevector(u32x8(x), u32x8(x), 1, 0, 3, 2, 5, 4, 7, 6));
    }
    [[gnu::target("avx2")]] inline u64x4 montgomery_reduce(u64x4 x, uint32_t mod, uint32_t imod) {
        auto x_ninv = u64x4(_mm256_mul_epu32(__m256i(x), __m256i() + imod));
        x += u64x4(_mm256_mul_epu32(__m256i(x_ninv), __m256i() + mod));
        return swap_bytes(x);
    }

    [[gnu::target("avx2")]] inline u64x4 montgomery_mul(u64x4 x, u64x4 y, uint32_t mod, uint32_t imod) {
        return montgomery_reduce(u64x4(_mm256_mul_epu32(__m256i(x), __m256i(y))), mod, imod);
    }
    [[gnu::target("avx2")]] inline u32x8 montgomery_mul(u32x8 x, u32x8 y, uint32_t mod, uint32_t imod) {
        return u32x8(montgomery_mul(u64x4(x), u64x4(y), mod, imod)) |
               u32x8(swap_bytes(montgomery_mul(u64x4(swap_bytes(x)), u64x4(swap_bytes(y)), mod, imod)));
    }
    [[gnu::target("avx2")]] inline dx4 rotate_right(dx4 x) {
        static constexpr u64x4 shuffler = {3, 0, 1, 2};
        return __builtin_shuffle(x, shuffler);
    }

    template<std::size_t Align = 32>
    [[gnu::target("avx2")]] inline bool is_aligned(const auto* p) noexcept {
        return (reinterpret_cast<std::uintptr_t>(p) % Align) == 0;
    }

    template<class Target>
    [[gnu::target("avx2")]] inline Target& vector_cast(auto &&p) {
        return *reinterpret_cast<Target*>(std::assume_aligned<alignof(Target)>(&p));
    }
}
#endif // CP_ALGO_UTIL_SIMD_HPP

#line 1 "cp-algo/util/simd.hpp"


#include <experimental/simd>
#include <cstdint>
#include <cstddef>
#include <memory>
namespace cp_algo {
    template<typename T, size_t len>
    using simd [[gnu::vector_size(len * sizeof(T))]] = T;
    using i64x4 = simd<int64_t, 4>;
    using u64x4 = simd<uint64_t, 4>;
    using u32x8 = simd<uint32_t, 8>;
    using i32x4 = simd<int32_t, 4>;
    using u32x4 = simd<uint32_t, 4>;
    using i16x4 = simd<int16_t, 4>;
    using u8x32 = simd<uint8_t, 32>;
    using dx4 = simd<double, 4>;

    [[gnu::target("avx2")]] inline dx4 abs(dx4 a) {
    return a < 0 ? -a : a;
    }

    // https://stackoverflow.com/a/77376595
    // works for ints in (-2^51, 2^51)
    static constexpr dx4 magic = dx4() + (3ULL << 51);
    [[gnu::target("avx2")]] inline i64x4 lround(dx4 x) {
        return i64x4(x + magic) - i64x4(magic);
    }
    [[gnu::target("avx2")]] inline dx4 to_double(i64x4 x) {
        return dx4(x + i64x4(magic)) - magic;
    }

    [[gnu::target("avx2")]] inline dx4 round(dx4 a) {
        return dx4{
            std::nearbyint(a[0]),
            std::nearbyint(a[1]),
            std::nearbyint(a[2]),
            std::nearbyint(a[3])
        };
    }

    [[gnu::target("avx2")]] inline u64x4 low32(u64x4 x) {
        return x & uint32_t(-1);
    }
    [[gnu::target("avx2")]] inline auto swap_bytes(auto x) {
        return decltype(x)(__builtin_shufflevector(u32x8(x), u32x8(x), 1, 0, 3, 2, 5, 4, 7, 6));
    }
    [[gnu::target("avx2")]] inline u64x4 montgomery_reduce(u64x4 x, uint32_t mod, uint32_t imod) {
        auto x_ninv = u64x4(_mm256_mul_epu32(__m256i(x), __m256i() + imod));
        x += u64x4(_mm256_mul_epu32(__m256i(x_ninv), __m256i() + mod));
        return swap_bytes(x);
    }

    [[gnu::target("avx2")]] inline u64x4 montgomery_mul(u64x4 x, u64x4 y, uint32_t mod, uint32_t imod) {
        return montgomery_reduce(u64x4(_mm256_mul_epu32(__m256i(x), __m256i(y))), mod, imod);
    }
    [[gnu::target("avx2")]] inline u32x8 montgomery_mul(u32x8 x, u32x8 y, uint32_t mod, uint32_t imod) {
        return u32x8(montgomery_mul(u64x4(x), u64x4(y), mod, imod)) |
               u32x8(swap_bytes(montgomery_mul(u64x4(swap_bytes(x)), u64x4(swap_bytes(y)), mod, imod)));
    }
    [[gnu::target("avx2")]] inline dx4 rotate_right(dx4 x) {
        static constexpr u64x4 shuffler = {3, 0, 1, 2};
        return __builtin_shuffle(x, shuffler);
    }

    template<std::size_t Align = 32>
    [[gnu::target("avx2")]] inline bool is_aligned(const auto* p) noexcept {
        return (reinterpret_cast<std::uintptr_t>(p) % Align) == 0;
    }

    template<class Target>
    [[gnu::target("avx2")]] inline Target& vector_cast(auto &&p) {
        return *reinterpret_cast<Target*>(std::assume_aligned<alignof(Target)>(&p));
    }
}

Back to top page