<commit_before>#include <cstdint>
#include <cassert>

#include <immintrin.h>
#include <x86intrin.h>

#define packed_dword(x) _mm512_set1_epi32(x)
#define packed_byte(x) packed_dword((x) | ((x) << 8) | ((x) << 16) | ((x) << 24))

namespace base64 {

    namespace avx512 {

        namespace precalc {

            static uint32_t lookup_0[256];
            static uint32_t lookup_1[256];
            static uint32_t lookup_2[256];
            static uint32_t lookup_3[256];
        }

        void initalize_lookup() {

            using namespace precalc;

            for (int i=0; i < 256; i++) {
                lookup_0[i] = 0x80000000;
                lookup_1[i] = 0x80000000;
                lookup_2[i] = 0x80000000;
                lookup_3[i] = 0x80000000;
            }

            for (int i=0; i < 64; i++) {
                const uint32_t val = static_cast<uint8_t>(base64::lookup[i]);

                lookup_0[val] = i;
                lookup_1[val] = i << 1*6;
                lookup_2[val] = i << 2*6;
                lookup_3[val] = i << 3*6;
            }
        }


        __m512i lookup_gather(const __m512i input, __mmask16& error) {

            const __m512i b0 = _mm512_and_si512(input, packed_dword(0x000000ff));
            const __m512i b1 = _mm512_and_si512(_mm512_srli_epi32(input, 1*8), packed_dword(0x000000ff));
            const __m512i b2 = _mm512_and_si512(_mm512_srli_epi32(input, 2*8), packed_dword(0x000000ff));
            const __m512i b3 = _mm512_and_si512(_mm512_srli_epi32(input, 3*8), packed_dword(0x000000ff));

            // do lookup
            const __m512i r0 = _mm512_i32gather_epi32(b0, (const int*)precalc::lookup_0, 4);
            const __m512i r1 = _mm512_i32gather_epi32(b1, (const int*)precalc::lookup_1, 4);
            const __m512i r2 = _mm512_i32gather_epi32(b2, (const int*)precalc::lookup_2, 4);
            const __m512i r3 = _mm512_i32gather_epi32(b3, (const int*)precalc::lookup_3, 4);
            const __m512i translated = _mm512_or_si512(r0,
                                       _mm512_or_si512(r1,
                                       _mm512_or_si512(r2, r3)));

            error = _mm512_cmplt_epi32_mask(translated, _mm512_set1_epi32(0));

            return translated;
        }


        // returns packed (a[i] >= lo and a[i] <= hi) ? 0x7f : 0x00;
        // assertion a[i] < 0x80
        __m512i _mm512_range_mask_7bit(const __m512i a, const __m512i lo, const __m512i hi) {

            const __m512i L = _mm512_add_epi32(a, lo);
            const __m512i H = _mm512_add_epi32(a, hi);

            /*
                MSB  H   L  |  R
                 0   0   0  |  0
                 0   0   1  |  0
                 0   1   0  |  0
                 0   1   1  |  0
                 1   0   0  |  0
                 1   0   1  |  1<commit_msg>decode: AVX512 gather-based proc slightly faster<commit_after>#include <cstdint>
#include <cassert>

#include <immintrin.h>
#include <x86intrin.h>

#define packed_dword(x) _mm512_set1_epi32(x)
#define packed_byte(x) packed_dword((x) | ((x) << 8) | ((x) << 16) | ((x) << 24))

namespace base64 {

    namespace avx512 {

        namespace precalc {

            static uint32_t lookup_0[256];
            static uint32_t lookup_1[256];
            static uint32_t lookup_2[256];
            static uint32_t lookup_3[256];
        }

        void initalize_lookup() {

            using namespace precalc;

            for (int i=0; i < 256; i++) {
                lookup_0[i] = 0x80000000;
                lookup_1[i] = 0x80000000;
                lookup_2[i] = 0x80000000;
                lookup_3[i] = 0x80000000;
            }

            for (int i=0; i < 64; i++) {
                const uint32_t val = static_cast<uint8_t>(base64::lookup[i]);

                lookup_0[val] = i;
                lookup_1[val] = i << 1*6;
                lookup_2[val] = i << 2*6;
                lookup_3[val] = i << 3*6;
            }
        }


        const uint8_t OR_ALL = 0xfe;


        __m512i lookup_gather(const __m512i input, __mmask16& error) {

            const __m512i b0 = _mm512_and_si512(input, packed_dword(0x000000ff));
            const __m512i b1 = _mm512_and_si512(_mm512_srli_epi32(input, 1*8), packed_dword(0x000000ff));
            const __m512i b2 = _mm512_and_si512(_mm512_srli_epi32(input, 2*8), packed_dword(0x000000ff));
            const __m512i b3 = _mm512_srli_epi32(input, 3*8);

            // do lookup
            const __m512i r0 = _mm512_i32gather_epi32(b0, (const int*)precalc::lookup_0, 4);
            const __m512i r1 = _mm512_i32gather_epi32(b1, (const int*)precalc::lookup_1, 4);
            const __m512i r2 = _mm512_i32gather_epi32(b2, (const int*)precalc::lookup_2, 4);
            const __m512i r3 = _mm512_i32gather_epi32(b3, (const int*)precalc::lookup_3, 4);

            // r0 | r1 | r2 | r3
            const __m512i translated = _mm512_or_si512(r0, _mm512_ternarylogic_epi32(r1, r2, r3, OR_ALL));

            error = _mm512_cmplt_epi32_mask(translated, _mm512_set1_epi32(0));

            return translated;
        }


        // returns packed (a[i] >= lo and a[i] <= hi) ? 0x7f : 0x00;
        // assertion a[i] < 0x80
        __m512i _mm512_range_mask_7bit(const __m512i a, const __m512i lo, const __m512i hi) {

            const __m512i L = _mm512_add_epi32(a, lo);
            const __m512i H = _mm512_add_epi32(a, hi);

            /*
                MSB  H   L  |  R
                 0   0   0  |  0
                 0   0   1  |  0
                 0   1   0  |  0
                 0   1   1  |  0
                 1   0   0  |  0
                 1   0   1  |  1
{'fraction_non_alphanumeric': 0.10560764196002123, 'fraction_numerical': 0.1247125420130904, 'mean_word_length': 1.751338199513382, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 8, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6925494', 'n_tokens_mistral': 2524, 'n_tokens_neox': 2046, 'n_words': 430}