libdpf/include/dpf/packed_lane_arithmetic.hpp

/// @file dpf/packed_lane_arithmetic.hpp
/// @brief AVX2 lane arithmetic for 2-bit (Z/4Z) and 4-bit (Z/16Z) leaves.
///
/// Vector-vector add of 2-bit lanes is the bitsliced form measured fastest
/// on AVX2 (carry stays in bit 0 of each pair). 4-bit add isolates nibbles
/// inside `add_epi8` so a sum cannot spill into the next nibble; that is
/// shorter than an unpack/pack into 16-bit lanes.
///
/// Scalar-vector multiply builds a 16-entry `pshufb` table, which won the
/// scalar-vector timings. Vector-vector multiply is bitsliced mod 4, and
/// `mullo_epi16` on nibbles split into even and odd bytes mod 16.

#ifndef LIBDPF_INCLUDE_DPF_PACKED_LANE_ARITHMETIC_HPP__
#define LIBDPF_INCLUDE_DPF_PACKED_LANE_ARITHMETIC_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <type_traits>

#include "hedley/hedley.h"
#include "simde/simde/x86/avx2.h"

#include "dpf/nyble.hpp"
#include "dpf/twobit.hpp"

namespace dpf
{
namespace lane_arith
{

namespace detail
{

template <unsigned Bits>
HEDLEY_ALWAYS_INLINE
unsigned lane_at(unsigned byte, unsigned shift) noexcept
{
    constexpr unsigned mask = (1u << Bits) - 1u;
    return (byte >> shift) & mask;
}

template <unsigned Bits, typename NodeT, typename Op>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
NodeT apply_bytes(const NodeT & a, const NodeT & b, Op op) noexcept
{
    static_assert(Bits == 2 || Bits == 4);
    constexpr unsigned mask = (1u << Bits) - 1u;
    constexpr unsigned per_byte = 8u / Bits;
    NodeT out{};
    const auto * ab = reinterpret_cast<const unsigned char *>(std::addressof(a));
    const auto * bb = reinterpret_cast<const unsigned char *>(std::addressof(b));
    auto * cb = reinterpret_cast<unsigned char *>(std::addressof(out));
    for (std::size_t i = 0; i < sizeof(NodeT); ++i)
    {
        unsigned packed = 0;
        for (unsigned s = 0; s < per_byte; ++s)
        {
            const unsigned shift = s * Bits;
            const unsigned av = lane_at<Bits>(ab[i], shift);
            const unsigned bv = lane_at<Bits>(bb[i], shift);
            packed |= (op(av, bv) & mask) << shift;
        }
        cb[i] = static_cast<unsigned char>(packed);
    }
    return out;
}

inline simde__m128i nibble_lut(const unsigned char lut[16]) noexcept
{
    simde__m128i table;
    std::memcpy(&table, lut, sizeof(table));
    return table;
}

inline simde__m256i nibble_lut256(const unsigned char lut[16]) noexcept
{
    alignas(32) unsigned char both[32];
    std::memcpy(both, lut, 16);
    std::memcpy(both + 16, lut, 16);
    simde__m256i table;
    std::memcpy(&table, both, sizeof(table));
    return table;
}

inline void fill_epi2_mul_lut(unsigned k, unsigned char lut[16]) noexcept
{
    k &= 3u;
    for (unsigned n = 0; n < 16u; ++n)
    {
        const unsigned l0 = ((n & 3u) * k) & 3u;
        const unsigned l1 = (((n >> 2) & 3u) * k) & 3u;
        lut[n] = static_cast<unsigned char>(l0 | (l1 << 2));
    }
}

inline void fill_epi4_mul_lut(unsigned k, unsigned char lut[16]) noexcept
{
    k &= 0x0fu;
    for (unsigned n = 0; n < 16u; ++n)
    {
        lut[n] = static_cast<unsigned char>((n * k) & 0x0fu);
    }
}

inline simde__m128i shuffle_nibbles(simde__m128i table, simde__m128i a) noexcept
{
    const auto m = simde_mm_set1_epi8(0x0f);
    const auto lo = simde_mm_shuffle_epi8(table, simde_mm_and_si128(a, m));
    const auto hi_idx = simde_mm_and_si128(simde_mm_srli_epi16(a, 4), m);
    const auto hi = simde_mm_shuffle_epi8(table, hi_idx);
    return simde_mm_or_si128(simde_mm_and_si128(lo, m),
        simde_mm_slli_epi16(simde_mm_and_si128(hi, m), 4));
}

inline simde__m256i shuffle_nibbles(simde__m256i table, simde__m256i a) noexcept
{
    const auto m = simde_mm256_set1_epi8(0x0f);
    const auto lo = simde_mm256_shuffle_epi8(table, simde_mm256_and_si256(a, m));
    const auto hi_idx = simde_mm256_and_si256(simde_mm256_srli_epi16(a, 4), m);
    const auto hi = simde_mm256_shuffle_epi8(table, hi_idx);
    return simde_mm256_or_si256(simde_mm256_and_si256(lo, m),
        simde_mm256_slli_epi16(simde_mm256_and_si256(hi, m), 4));
}

/// Low nibble of every byte, product mod 16. Even and odd bytes are split
/// so a product in one byte cannot land in the next.
inline simde__m128i mul_low_nibbles(simde__m128i a, simde__m128i b) noexcept
{
    const auto lane = simde_mm_set1_epi16(0x000f);
    const auto ae = simde_mm_and_si128(a, lane);
    const auto be = simde_mm_and_si128(b, lane);
    const auto pe = simde_mm_and_si128(simde_mm_mullo_epi16(ae, be), lane);
    const auto ao = simde_mm_and_si128(simde_mm_srli_epi16(a, 8), lane);
    const auto bo = simde_mm_and_si128(simde_mm_srli_epi16(b, 8), lane);
    const auto po = simde_mm_and_si128(simde_mm_mullo_epi16(ao, bo), lane);
    return simde_mm_or_si128(pe, simde_mm_slli_epi16(po, 8));
}

inline simde__m256i mul_low_nibbles(simde__m256i a, simde__m256i b) noexcept
{
    const auto lane = simde_mm256_set1_epi16(0x000f);
    const auto ae = simde_mm256_and_si256(a, lane);
    const auto be = simde_mm256_and_si256(b, lane);
    const auto pe = simde_mm256_and_si256(simde_mm256_mullo_epi16(ae, be), lane);
    const auto ao = simde_mm256_and_si256(simde_mm256_srli_epi16(a, 8), lane);
    const auto bo = simde_mm256_and_si256(simde_mm256_srli_epi16(b, 8), lane);
    const auto po = simde_mm256_and_si256(simde_mm256_mullo_epi16(ao, bo), lane);
    return simde_mm256_or_si256(pe, simde_mm256_slli_epi16(po, 8));
}

}  // namespace detail

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT add_mod2(const NodeT & a, const NodeT & b) noexcept
{
    return detail::apply_bytes<2>(a, b, [](unsigned x, unsigned y) {
        return (x + y) & 3u;
    });
}

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT sub_mod2(const NodeT & a, const NodeT & b) noexcept
{
    return detail::apply_bytes<2>(a, b, [](unsigned x, unsigned y) {
        return (x - y) & 3u;
    });
}

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT mul_mod2(const NodeT & a, dpf::twobit k) noexcept
{
    const unsigned s = static_cast<unsigned>(k);
    return detail::apply_bytes<2>(a, a, [s](unsigned x, unsigned) {
        return (x * s) & 3u;
    });
}

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT mullo_mod2(const NodeT & a, const NodeT & b) noexcept
{
    return detail::apply_bytes<2>(a, b, [](unsigned x, unsigned y) {
        return (x * y) & 3u;
    });
}

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT add_mod4(const NodeT & a, const NodeT & b) noexcept
{
    return detail::apply_bytes<4>(a, b, [](unsigned x, unsigned y) {
        return (x + y) & 0x0fu;
    });
}

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT sub_mod4(const NodeT & a, const NodeT & b) noexcept
{
    return detail::apply_bytes<4>(a, b, [](unsigned x, unsigned y) {
        return (x - y) & 0x0fu;
    });
}

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT mul_mod4(const NodeT & a, dpf::nyble k) noexcept
{
    const unsigned s = static_cast<unsigned>(k);
    return detail::apply_bytes<4>(a, a, [s](unsigned x, unsigned) {
        return (x * s) & 0x0fu;
    });
}

template <typename NodeT>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
NodeT mullo_mod4(const NodeT & a, const NodeT & b) noexcept
{
    return detail::apply_bytes<4>(a, b, [](unsigned x, unsigned y) {
        return (x * y) & 0x0fu;
    });
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i add_epi2(simde__m128i a, simde__m128i b) noexcept
{
    const auto m = simde_mm_set1_epi8(0x55);
    const auto a0 = simde_mm_and_si128(a, m);
    const auto a1 = simde_mm_and_si128(simde_mm_srli_epi16(a, 1), m);
    const auto b0 = simde_mm_and_si128(b, m);
    const auto b1 = simde_mm_and_si128(simde_mm_srli_epi16(b, 1), m);
    const auto bit0 = simde_mm_xor_si128(a0, b0);
    const auto carry = simde_mm_and_si128(a0, b0);
    const auto bit1 = simde_mm_xor_si128(simde_mm_xor_si128(a1, b1), carry);
    return simde_mm_or_si128(bit0, simde_mm_slli_epi16(bit1, 1));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i add_epi2(simde__m256i a, simde__m256i b) noexcept
{
    const auto m = simde_mm256_set1_epi8(0x55);
    const auto a0 = simde_mm256_and_si256(a, m);
    const auto a1 = simde_mm256_and_si256(simde_mm256_srli_epi16(a, 1), m);
    const auto b0 = simde_mm256_and_si256(b, m);
    const auto b1 = simde_mm256_and_si256(simde_mm256_srli_epi16(b, 1), m);
    const auto bit0 = simde_mm256_xor_si256(a0, b0);
    const auto carry = simde_mm256_and_si256(a0, b0);
    const auto bit1 = simde_mm256_xor_si256(simde_mm256_xor_si256(a1, b1), carry);
    return simde_mm256_or_si256(bit0, simde_mm256_slli_epi16(bit1, 1));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i sub_epi2(simde__m128i a, simde__m128i b) noexcept
{
    const auto m = simde_mm_set1_epi8(0x55);
    const auto a0 = simde_mm_and_si128(a, m);
    const auto a1 = simde_mm_and_si128(simde_mm_srli_epi16(a, 1), m);
    const auto b0 = simde_mm_and_si128(b, m);
    const auto b1 = simde_mm_and_si128(simde_mm_srli_epi16(b, 1), m);
    const auto bit0 = simde_mm_xor_si128(a0, b0);
    const auto borrow = simde_mm_andnot_si128(a0, b0);
    const auto bit1 = simde_mm_xor_si128(simde_mm_xor_si128(a1, b1), borrow);
    return simde_mm_or_si128(bit0, simde_mm_slli_epi16(bit1, 1));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i sub_epi2(simde__m256i a, simde__m256i b) noexcept
{
    const auto m = simde_mm256_set1_epi8(0x55);
    const auto a0 = simde_mm256_and_si256(a, m);
    const auto a1 = simde_mm256_and_si256(simde_mm256_srli_epi16(a, 1), m);
    const auto b0 = simde_mm256_and_si256(b, m);
    const auto b1 = simde_mm256_and_si256(simde_mm256_srli_epi16(b, 1), m);
    const auto bit0 = simde_mm256_xor_si256(a0, b0);
    const auto borrow = simde_mm256_andnot_si256(a0, b0);
    const auto bit1 = simde_mm256_xor_si256(simde_mm256_xor_si256(a1, b1), borrow);
    return simde_mm256_or_si256(bit0, simde_mm256_slli_epi16(bit1, 1));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i mullo_epi2(simde__m128i a, simde__m128i b) noexcept
{
    const auto m = simde_mm_set1_epi8(0x55);
    const auto a0 = simde_mm_and_si128(a, m);
    const auto a1 = simde_mm_and_si128(simde_mm_srli_epi16(a, 1), m);
    const auto b0 = simde_mm_and_si128(b, m);
    const auto b1 = simde_mm_and_si128(simde_mm_srli_epi16(b, 1), m);
    const auto bit0 = simde_mm_and_si128(a0, b0);
    const auto bit1 = simde_mm_xor_si128(
        simde_mm_and_si128(a0, b1), simde_mm_and_si128(a1, b0));
    return simde_mm_or_si128(bit0, simde_mm_slli_epi16(bit1, 1));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i mullo_epi2(simde__m256i a, simde__m256i b) noexcept
{
    const auto m = simde_mm256_set1_epi8(0x55);
    const auto a0 = simde_mm256_and_si256(a, m);
    const auto a1 = simde_mm256_and_si256(simde_mm256_srli_epi16(a, 1), m);
    const auto b0 = simde_mm256_and_si256(b, m);
    const auto b1 = simde_mm256_and_si256(simde_mm256_srli_epi16(b, 1), m);
    const auto bit0 = simde_mm256_and_si256(a0, b0);
    const auto bit1 = simde_mm256_xor_si256(
        simde_mm256_and_si256(a0, b1), simde_mm256_and_si256(a1, b0));
    return simde_mm256_or_si256(bit0, simde_mm256_slli_epi16(bit1, 1));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i mul_epi2(simde__m128i a, dpf::twobit k) noexcept
{
    unsigned char lut[16];
    detail::fill_epi2_mul_lut(static_cast<unsigned>(k), lut);
    return detail::shuffle_nibbles(detail::nibble_lut(lut), a);
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i mul_epi2(simde__m256i a, dpf::twobit k) noexcept
{
    unsigned char lut[16];
    detail::fill_epi2_mul_lut(static_cast<unsigned>(k), lut);
    return detail::shuffle_nibbles(detail::nibble_lut256(lut), a);
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i add_epi4(simde__m128i a, simde__m128i b) noexcept
{
    const auto lo_m = simde_mm_set1_epi8(0x0f);
    const auto hi_m = simde_mm_set1_epi8(static_cast<int8_t>(0xf0));
    const auto lo = simde_mm_add_epi8(
        simde_mm_and_si128(a, lo_m), simde_mm_and_si128(b, lo_m));
    const auto hi = simde_mm_add_epi8(
        simde_mm_and_si128(a, hi_m), simde_mm_and_si128(b, hi_m));
    return simde_mm_or_si128(
        simde_mm_and_si128(lo, lo_m), simde_mm_and_si128(hi, hi_m));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i add_epi4(simde__m256i a, simde__m256i b) noexcept
{
    const auto lo_m = simde_mm256_set1_epi8(0x0f);
    const auto hi_m = simde_mm256_set1_epi8(static_cast<int8_t>(0xf0));
    const auto lo = simde_mm256_add_epi8(
        simde_mm256_and_si256(a, lo_m), simde_mm256_and_si256(b, lo_m));
    const auto hi = simde_mm256_add_epi8(
        simde_mm256_and_si256(a, hi_m), simde_mm256_and_si256(b, hi_m));
    return simde_mm256_or_si256(
        simde_mm256_and_si256(lo, lo_m), simde_mm256_and_si256(hi, hi_m));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i sub_epi4(simde__m128i a, simde__m128i b) noexcept
{
    const auto lo_m = simde_mm_set1_epi8(0x0f);
    const auto hi_m = simde_mm_set1_epi8(static_cast<int8_t>(0xf0));
    const auto lo = simde_mm_sub_epi8(
        simde_mm_and_si128(a, lo_m), simde_mm_and_si128(b, lo_m));
    const auto hi = simde_mm_sub_epi8(
        simde_mm_and_si128(a, hi_m), simde_mm_and_si128(b, hi_m));
    return simde_mm_or_si128(
        simde_mm_and_si128(lo, lo_m), simde_mm_and_si128(hi, hi_m));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i sub_epi4(simde__m256i a, simde__m256i b) noexcept
{
    const auto lo_m = simde_mm256_set1_epi8(0x0f);
    const auto hi_m = simde_mm256_set1_epi8(static_cast<int8_t>(0xf0));
    const auto lo = simde_mm256_sub_epi8(
        simde_mm256_and_si256(a, lo_m), simde_mm256_and_si256(b, lo_m));
    const auto hi = simde_mm256_sub_epi8(
        simde_mm256_and_si256(a, hi_m), simde_mm256_and_si256(b, hi_m));
    return simde_mm256_or_si256(
        simde_mm256_and_si256(lo, lo_m), simde_mm256_and_si256(hi, hi_m));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i mullo_epi4(simde__m128i a, simde__m128i b) noexcept
{
    const auto m = simde_mm_set1_epi8(0x0f);
    const auto lo = detail::mul_low_nibbles(a, b);
    const auto ah = simde_mm_and_si128(simde_mm_srli_epi16(a, 4), m);
    const auto bh = simde_mm_and_si128(simde_mm_srli_epi16(b, 4), m);
    const auto hi = detail::mul_low_nibbles(ah, bh);
    return simde_mm_or_si128(lo, simde_mm_slli_epi16(hi, 4));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i mullo_epi4(simde__m256i a, simde__m256i b) noexcept
{
    const auto m = simde_mm256_set1_epi8(0x0f);
    const auto lo = detail::mul_low_nibbles(a, b);
    const auto ah = simde_mm256_and_si256(simde_mm256_srli_epi16(a, 4), m);
    const auto bh = simde_mm256_and_si256(simde_mm256_srli_epi16(b, 4), m);
    const auto hi = detail::mul_low_nibbles(ah, bh);
    return simde_mm256_or_si256(lo, simde_mm256_slli_epi16(hi, 4));
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m128i mul_epi4(simde__m128i a, dpf::nyble k) noexcept
{
    unsigned char lut[16];
    detail::fill_epi4_mul_lut(static_cast<unsigned>(k), lut);
    return detail::shuffle_nibbles(detail::nibble_lut(lut), a);
}

HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
HEDLEY_PURE
simde__m256i mul_epi4(simde__m256i a, dpf::nyble k) noexcept
{
    unsigned char lut[16];
    detail::fill_epi4_mul_lut(static_cast<unsigned>(k), lut);
    return detail::shuffle_nibbles(detail::nibble_lut256(lut), a);
}

}  // namespace lane_arith
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PACKED_LANE_ARITHMETIC_HPP__