ascii.cpp

#include <cstdio>
#include <cstdint>
#include <cstring>
#include <cassert>
#include <vector>
#include <algorithm>

#include <sys/time.h>

static inline int ascii_std(const uint8_t *data, int len)
{
    return !std::any_of(data, data+len, [] (int8_t b) { return b < 0; });
}

static inline int ascii_u64(const uint8_t *data, int len)
{
    uint8_t orall = 0;

    if (len >= 16) {

        uint64_t or1 = 0, or2 = 0;
        const uint8_t *data2 = data+8;

        do {
            or1 |= *(const uint64_t *)data;
            or2 |= *(const uint64_t *)data2;
            data += 16;
            data2 += 16;
            len -= 16;
        } while (len >= 16);

        /*
         * Idea from Benny Halevy <bhalevy@scylladb.com>
         * - 7-th bit set   ==> orall = !(non-zero) - 1 = 0 - 1 = 0xFF
         * - 7-th bit clear ==> orall = !0 - 1          = 1 - 1 = 0x00
         */
        orall = !((or1 | or2) & 0x8080808080808080ULL) - 1;
    }

    while (len--)
        orall |= *data++;

    return orall < 0x80;
}

#if defined(__x86_64__)
#include <x86intrin.h>

static inline int ascii_simd(const uint8_t *data, int len)
{
    if (len >= 32) {
        const uint8_t *data2 = data+16;

        __m128i or1 = _mm_set1_epi8(0), or2 = or1;

        while (len >= 32) {
            __m128i input1 = _mm_loadu_si128((const __m128i *)data);
            __m128i input2 = _mm_loadu_si128((const __m128i *)data2);

            or1 = _mm_or_si128(or1, input1);
            or2 = _mm_or_si128(or2, input2);

            data += 32;
            data2 += 32;
            len -= 32;
        }

        or1 = _mm_or_si128(or1, or2);
        if (_mm_movemask_epi8(_mm_cmplt_epi8(or1, _mm_set1_epi8(0))))
            return 0;
    }

    return ascii_u64(data, len);
}

#elif defined(__aarch64__)
#include <arm_neon.h>

static inline int ascii_simd(const uint8_t *data, int len)
{
    if (len >= 32) {
        const uint8_t *data2 = data+16;

        uint8x16_t or1 = vdupq_n_u8(0), or2 = or1;

        while (len >= 32) {
            const uint8x16_t input1 = vld1q_u8(data);
            const uint8x16_t input2 = vld1q_u8(data2);

            or1 = vorrq_u8(or1, input1);
            or2 = vorrq_u8(or2, input2);

            data += 32;
            data2 += 32;
            len -= 32;
        }

        or1 = vorrq_u8(or1, or2);
        if (vmaxvq_u8(or1) >= 0x80)
            return 0;
    }

    return ascii_u64(data, len);
}

#endif

struct ftab {
    const char *name;
    int (*func)(const uint8_t *data, int len);
};

static const std::vector<ftab> _f = {
    {
        .name = "std",
        .func = ascii_std,
    }, {
        .name = "u64",
        .func = ascii_u64,
    }, {
        .name = "simd",
        .func = ascii_simd,
    },
};

static void load_test_buf(uint8_t *data, int len)
{
    uint8_t v = 0;

    for (int i = 0; i < len; ++i) {
        data[i] = v++;
        v &= 0x7F;
    }
}

static void bench(const struct ftab &f, const uint8_t *data, int len)
{
    const int loops = 1024*1024*1024/len;
    int ret = 1;
    double time_aligned, time_unaligned, size;
    struct timeval tv1, tv2;

    fprintf(stderr, "bench %s (%d bytes)... ", f.name, len);

    /* aligned */
    gettimeofday(&tv1, 0);
    for (int i = 0; i < loops; ++i)
        ret &= f.func(data, len);
    gettimeofday(&tv2, 0);
    time_aligned = tv2.tv_usec - tv1.tv_usec;
    time_aligned = time_aligned / 1000000 + tv2.tv_sec - tv1.tv_sec;

    /* unaligned */
    gettimeofday(&tv1, 0);
    for (int i = 0; i < loops; ++i)
        ret &= f.func(data+1, len);
    gettimeofday(&tv2, 0);
    time_unaligned = tv2.tv_usec - tv1.tv_usec;
    time_unaligned = time_unaligned / 1000000 + tv2.tv_sec - tv1.tv_sec;

    printf("%s ", ret?"pass":"FAIL");

    size = ((double)len * loops) / (1024*1024);
    printf("%.0f/%.0f MB/s\n", size / time_aligned, size / time_unaligned);
}

static void test(const struct ftab &f, uint8_t *data, int len)
{
    int error = 0;

    fprintf(stderr, "test %s (%d bytes)... ", f.name, len);

    /* positive */
    error |= !f.func(data, len);

    /* negative */
    if (len < 100*1024) {
        for (int i = 0; i < len; ++i) {
            data[i] += 0x80;
            error |= f.func(data, len);
            data[i] -= 0x80;
        }
    }

    printf("%s\n", error ? "FAIL" : "pass");
}

/* ./ascii [test|bench] [alg] */
int main(int argc, const char *argv[])
{
    int do_test = 1, do_bench = 1;
    const char *alg = NULL;

    if (argc > 1) {
        do_bench &= !!strcmp(argv[1], "test");
        do_test &= !!strcmp(argv[1], "bench");
    }

    if (do_bench && argc > 2)
        alg = argv[2];

    const std::vector<int> size = { 9, 16+1, 32-1, 128+1, 1024+15,
                                    16*1024+1, 64*1024+15, 1024*1024 };

    int max_size = *std::max_element(size.begin(), size.end());
    uint8_t *_data = new uint8_t[max_size+1];
    assert(((uintptr_t)_data & 7) == 0);
    uint8_t *data = _data+1;    /* Unalign buffer address */

    _data[0] = 0;
    load_test_buf(data, max_size);

    if (do_test) {
        printf("==================== Test ====================\n");
        for (int sz : size) {
            for (auto &f : _f) {
                test(f, data, sz);
            }
        }
    }

    if (do_bench) {
        printf("==================== Bench ====================\n");
        for (int sz : size) {
            for (auto &f : _f) {
                if (!alg || strcmp(alg, f.name) == 0)
                    bench(f, _data, sz);
            }
            printf("-----------------------------------------------\n");
        }
    }

    delete _data;
    return 0;
}