Project

General

Profile

Bug #19875 » bytecount.c

missing/bytecount.c - Freaky (Thomas Hurst), 09/19/2023 11:06 PM

 

#include <assert.h>
#include <stddef.h>
#include <stdint.h>

#if (defined(__i386__) || defined(__x86_64__))
#define HAVE_AVX2 1
#include <immintrin.h>

static const char AVX2_MASK[] = {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255
};

__attribute__((target("avx2")))
size_t
avx2_bytecount(const char *haystack, const char needle, size_t haystack_len) {
assert (haystack_len >= 32);

#define SUM_ADD(count, u8s, temp) do { \
temp = _mm256_sad_epu8(u8s, _mm256_setzero_si256()); \
count += _mm256_extract_epi64(temp, 0) + _mm256_extract_epi64(temp, 1) + \
_mm256_extract_epi64(temp, 2) + _mm256_extract_epi64(temp, 3); \
} while(0)

#define mm256_from_offset(slice, offset) _mm256_loadu_si256((__m256i*)(slice + offset))

size_t offset = 0;
size_t count = 0;

__m256i sums;
__m256i needles = _mm256_set1_epi8((char)needle);

// 8160
while (haystack_len >= offset + 32 * 255) {
__m256i counts = _mm256_setzero_si256();
for (int i = 0; i < 255; ++i) {
counts = _mm256_sub_epi8(
counts,
_mm256_cmpeq_epi8(mm256_from_offset(haystack, offset), needles)
);
offset += 32;
}
SUM_ADD(count, counts, sums);
}

// 4096
if (haystack_len >= offset + 32 * 128) {
__m256i counts = _mm256_setzero_si256();
for (int i = 0; i < 128; ++i) {
counts = _mm256_sub_epi8(
counts,
_mm256_cmpeq_epi8(mm256_from_offset(haystack, offset), needles)
);
offset += 32;
}
SUM_ADD(count, counts, sums);
}

// 32
__m256i counts = _mm256_setzero_si256();
for (size_t i = 0; i < (haystack_len - offset) / 32; ++i) {
counts = _mm256_sub_epi8(
counts,
_mm256_cmpeq_epi8(mm256_from_offset(haystack, offset + i * 32), needles)
);
}
if (haystack_len % 32 != 0) {
counts = _mm256_sub_epi8(
counts,
_mm256_and_si256(
_mm256_cmpeq_epi8(mm256_from_offset(haystack, haystack_len - 32), needles),
mm256_from_offset(AVX2_MASK, haystack_len % 32)
)
);
}
SUM_ADD(count, counts, sums);

return count;

#undef mm256_from_offset
#undef SUM_ADD
}

#define HAVE_SSE4_1
#include <smmintrin.h>

static const char SSE4_1_MASK[] = {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
};

__attribute__((target("sse4.1")))
size_t
sse41_bytecount(const char *haystack, const char needle, size_t haystack_len) {
assert(haystack_len >= 16);

#define mm_from_offset(haystack, offset) _mm_loadu_si128((__m128i *)(haystack + offset))
#define SUM_ADD(count, u8s, temp) do { \
temp = _mm_sad_epu8(u8s, _mm_setzero_si128()); \
count += _mm_extract_epi32(sums, 0) + _mm_extract_epi32(sums, 2); \
} while (0)

const char *ptr = (const char *)haystack;
size_t offset = 0;
size_t count = 0;

const __m128i needles = _mm_set1_epi8((char)needle);
__m128i sums;

// 4080
while (haystack_len >= offset + 16 * 255) {
__m128i counts = _mm_setzero_si128();

for (size_t i=0; i < 255; i++) {
counts = _mm_sub_epi8(
counts,
_mm_cmpeq_epi8(mm_from_offset(ptr, offset), needles)
);
offset += 16;
}

SUM_ADD(count, counts, sums);
}

// 2048
if (haystack_len >= offset + 16 * 128) {
__m128i counts = _mm_setzero_si128();

for (size_t i=0; i < 128; i++) {
counts = _mm_sub_epi8(
counts,
_mm_cmpeq_epi8(mm_from_offset(ptr, offset), needles)
);
offset += 16;
}

SUM_ADD(count, counts, sums);
}

// 16
__m128i counts = _mm_setzero_si128();
for (size_t i=0; i < ((haystack_len - offset) / 16); i++) {
counts = _mm_sub_epi8(
counts,
_mm_cmpeq_epi8(mm_from_offset(ptr, offset + i * 16), needles)
);
}

if (haystack_len % 16 != 0) {
counts = _mm_sub_epi8(
counts,
_mm_and_si128(
_mm_cmpeq_epi8(mm_from_offset(ptr, haystack_len - 16), needles),
mm_from_offset(SSE4_1_MASK, haystack_len % 16)
)
);
}
SUM_ADD(count, counts, sums);

return count;

#undef mm_from_offset
#undef SUM_ADD
}
#endif

static size_t
naive_bytecount(const char *haystack, const char needle, size_t haystack_len) {
size_t count = 0;
const char *ptr = haystack;
const char *end_ptr = ptr + haystack_len;

while (ptr < end_ptr) {
if (*ptr == needle) {
count++;
}
ptr++;
}

return count;
}

static size_t
naive_bytecount_32(const char *haystack, const char needle, size_t haystack_len) {
uint32_t count = 0;
const char *ptr = haystack;
const char *end_ptr = ptr + haystack_len;

while (ptr < end_ptr) {
if (*ptr == needle) {
count++;
}
ptr++;
}

return count;
}

static size_t
_fallback_bytecount(const char *haystack, const char needle, size_t haystack_len) {
if (haystack_len < UINT32_MAX) {
return naive_bytecount_32(haystack, needle, haystack_len);
}

return naive_bytecount(haystack, needle, haystack_len);
}

#ifdef HAVE_AVX2
__attribute__((target("sse4.1,avx2")))
static size_t
_avx2_bytecount(const char *haystack, const char needle, size_t haystack_len) {
if (haystack_len >= 32) {
return avx2_bytecount(haystack, needle, haystack_len);
}

#ifdef HAVE_SSE4_1
if (haystack_len >= 16) {
return sse41_bytecount(haystack, needle, haystack_len);
}
#endif

return _fallback_bytecount(haystack, needle, haystack_len);
}
#endif

#ifdef HAVE_SSE4_1
__attribute__((target("sse4.1")))
static size_t
_sse41_bytecount(const char *haystack, const char needle, size_t haystack_len) {
if (haystack_len >= 16) {
return sse41_bytecount(haystack, needle, haystack_len);
}

return _fallback_bytecount(haystack, needle, haystack_len);
}
#endif

__attribute__((unused))
static size_t
(*resolve_bytecount (void))(const char *, const char, size_t)
{
__builtin_cpu_init();
#ifdef HAVE_AVX2
if (__builtin_cpu_supports("avx2")) {
return _avx2_bytecount;
}
#endif

#ifdef HAVE_SSE4_1
if (__builtin_cpu_supports("sse4.1")) {
return _sse41_bytecount;
}
#endif

return _fallback_bytecount;
}

static size_t
inner_bytecount(const char *haystack, const char needle, size_t haystack_len)
__attribute__((ifunc("resolve_bytecount")));

static size_t
bytecount(const void *haystack, int needle, size_t haystack_len) {
return inner_bytecount((const char *)haystack, (const char)needle, haystack_len);
}
(5-5/5)