src/port/pg_popcount_x86.c - master - githubmirror/postgres

src/port/pg_popcount_x86.c
310 строк8 KB
Nathan Bossart
Suppress "has no symbols" linker warnings on macOS.
29 апр 2026, 20:25
29 апр 2026, 20:253dd42ee
/*-------------------------------------------------------------------------
 *
 * pg_popcount_x86.c
 *	  Holds the x86-64 pg_popcount() implementations.
 *
 * Copyright (c) 2024-2026, PostgreSQL Global Development Group
 *
 * IDENTIFICATION
 *	  src/port/pg_popcount_x86.c
 *
 *-------------------------------------------------------------------------
 */
#include "c.h"

#ifdef HAVE_X86_64_POPCNTQ

#ifdef USE_AVX512_POPCNT_WITH_RUNTIME_CHECK
#include <immintrin.h>
#endif

#include "port/pg_bitutils.h"
#include "port/pg_cpu.h"

/*
 * The SSE4.2 versions are built regardless of whether we are building the
 * AVX-512 versions.
 *
 * Technically, POPCNT is not part of SSE4.2, and isn't even a vector
 * operation, but in practice this is close enough, and "sse42" seems easier to
 * follow than "popcnt" for these names.
 */
static uint64 pg_popcount_sse42(const char *buf, int bytes);
static uint64 pg_popcount_masked_sse42(const char *buf, int bytes, uint8 mask);

/*
 * These are the AVX-512 implementations of the popcount functions.
 */
#ifdef USE_AVX512_POPCNT_WITH_RUNTIME_CHECK
static uint64 pg_popcount_avx512(const char *buf, int bytes);
static uint64 pg_popcount_masked_avx512(const char *buf, int bytes, uint8 mask);
#endif							/* USE_AVX512_POPCNT_WITH_RUNTIME_CHECK */

/*
 * The function pointers are initially set to "choose" functions.  These
 * functions will first set the pointers to the right implementations (base on
 * what the current CPU supports) and then will call the pointer to fulfill the
 * caller's request.
 */
static uint64 pg_popcount_choose(const char *buf, int bytes);
static uint64 pg_popcount_masked_choose(const char *buf, int bytes, uint8 mask);
uint64		(*pg_popcount_optimized) (const char *buf, int bytes) = pg_popcount_choose;
uint64		(*pg_popcount_masked_optimized) (const char *buf, int bytes, uint8 mask) = pg_popcount_masked_choose;


#ifdef USE_AVX512_POPCNT_WITH_RUNTIME_CHECK

/*
 * Returns true if the CPU supports the instructions required for the AVX-512
 * pg_popcount() implementation.
 */
static bool
pg_popcount_avx512_available(void)
{
	return x86_feature_available(PG_AVX512_BW) &&
		x86_feature_available(PG_AVX512_VPOPCNTDQ);
}

#endif							/* USE_AVX512_POPCNT_WITH_RUNTIME_CHECK */

/*
 * These functions get called on the first call to pg_popcount(), etc.
 * They detect whether we can use the asm implementations, and replace
 * the function pointers so that subsequent calls are routed directly to
 * the chosen implementation.
 */
static inline void
choose_popcount_functions(void)
{
	if (x86_feature_available(PG_POPCNT))
	{
		pg_popcount_optimized = pg_popcount_sse42;
		pg_popcount_masked_optimized = pg_popcount_masked_sse42;
	}
	else
	{
		pg_popcount_optimized = pg_popcount_portable;
		pg_popcount_masked_optimized = pg_popcount_masked_portable;
	}

#ifdef USE_AVX512_POPCNT_WITH_RUNTIME_CHECK
	if (pg_popcount_avx512_available())
	{
		pg_popcount_optimized = pg_popcount_avx512;
		pg_popcount_masked_optimized = pg_popcount_masked_avx512;
	}
#endif
}

static uint64
pg_popcount_choose(const char *buf, int bytes)
{
	choose_popcount_functions();
	return pg_popcount_optimized(buf, bytes);
}

static uint64
pg_popcount_masked_choose(const char *buf, int bytes, uint8 mask)
{
	choose_popcount_functions();
	return pg_popcount_masked(buf, bytes, mask);
}

#ifdef USE_AVX512_POPCNT_WITH_RUNTIME_CHECK

/*
 * pg_popcount_avx512
 *		Returns the number of 1-bits in buf
 */
pg_attribute_target("avx512vpopcntdq,avx512bw")
static uint64
pg_popcount_avx512(const char *buf, int bytes)
{
	__m512i		val,
				cnt;
	__m512i		accum = _mm512_setzero_si512();
	const char *final;
	int			tail_idx;
	__mmask64	mask = ~UINT64CONST(0);

	/*
	 * Align buffer down to avoid double load overhead from unaligned access.
	 * Calculate a mask to ignore preceding bytes.  Find start offset of final
	 * iteration and ensure it is not empty.
	 */
	mask <<= ((uintptr_t) buf) % sizeof(__m512i);
	tail_idx = (((uintptr_t) buf + bytes - 1) % sizeof(__m512i)) + 1;
	final = (const char *) TYPEALIGN_DOWN(sizeof(__m512i), buf + bytes - 1);
	buf = (const char *) TYPEALIGN_DOWN(sizeof(__m512i), buf);

	/*
	 * Iterate through all but the final iteration.  Starting from the second
	 * iteration, the mask is ignored.
	 */
	if (buf < final)
	{
		val = _mm512_maskz_loadu_epi8(mask, (const __m512i *) buf);
		cnt = _mm512_popcnt_epi64(val);
		accum = _mm512_add_epi64(accum, cnt);

		buf += sizeof(__m512i);
		mask = ~UINT64CONST(0);

		for (; buf < final; buf += sizeof(__m512i))
		{
			val = _mm512_load_si512((const __m512i *) buf);
			cnt = _mm512_popcnt_epi64(val);
			accum = _mm512_add_epi64(accum, cnt);
		}
	}

	/* Final iteration needs to ignore bytes that are not within the length */
	mask &= (~UINT64CONST(0) >> (sizeof(__m512i) - tail_idx));

	val = _mm512_maskz_loadu_epi8(mask, (const __m512i *) buf);
	cnt = _mm512_popcnt_epi64(val);
	accum = _mm512_add_epi64(accum, cnt);

	return _mm512_reduce_add_epi64(accum);
}

/*
 * pg_popcount_masked_avx512
 *		Returns the number of 1-bits in buf after applying the mask to each byte
 */
pg_attribute_target("avx512vpopcntdq,avx512bw")
static uint64
pg_popcount_masked_avx512(const char *buf, int bytes, uint8 mask)
{
	__m512i		val,
				vmasked,
				cnt;
	__m512i		accum = _mm512_setzero_si512();
	const char *final;
	int			tail_idx;
	__mmask64	bmask = ~UINT64CONST(0);
	const __m512i maskv = _mm512_set1_epi8(mask);

	/*
	 * Align buffer down to avoid double load overhead from unaligned access.
	 * Calculate a mask to ignore preceding bytes.  Find start offset of final
	 * iteration and ensure it is not empty.
	 */
	bmask <<= ((uintptr_t) buf) % sizeof(__m512i);
	tail_idx = (((uintptr_t) buf + bytes - 1) % sizeof(__m512i)) + 1;
	final = (const char *) TYPEALIGN_DOWN(sizeof(__m512i), buf + bytes - 1);
	buf = (const char *) TYPEALIGN_DOWN(sizeof(__m512i), buf);

	/*
	 * Iterate through all but the final iteration.  Starting from the second
	 * iteration, the mask is ignored.
	 */
	if (buf < final)
	{
		val = _mm512_maskz_loadu_epi8(bmask, (const __m512i *) buf);
		vmasked = _mm512_and_si512(val, maskv);
		cnt = _mm512_popcnt_epi64(vmasked);
		accum = _mm512_add_epi64(accum, cnt);

		buf += sizeof(__m512i);
		bmask = ~UINT64CONST(0);

		for (; buf < final; buf += sizeof(__m512i))
		{
			val = _mm512_load_si512((const __m512i *) buf);
			vmasked = _mm512_and_si512(val, maskv);
			cnt = _mm512_popcnt_epi64(vmasked);
			accum = _mm512_add_epi64(accum, cnt);
		}
	}

	/* Final iteration needs to ignore bytes that are not within the length */
	bmask &= (~UINT64CONST(0) >> (sizeof(__m512i) - tail_idx));

	val = _mm512_maskz_loadu_epi8(bmask, (const __m512i *) buf);
	vmasked = _mm512_and_si512(val, maskv);
	cnt = _mm512_popcnt_epi64(vmasked);
	accum = _mm512_add_epi64(accum, cnt);

	return _mm512_reduce_add_epi64(accum);
}

#endif							/* USE_AVX512_POPCNT_WITH_RUNTIME_CHECK */

/*
 * pg_popcount64_sse42
 *		Return the number of 1 bits set in word
 */
static inline int
pg_popcount64_sse42(uint64 word)
{
#ifdef _MSC_VER
	return __popcnt64(word);
#else
	uint64		res;

__asm__ __volatile__(" popcntq %1,%0\n":"=q"(res):"rm"(word):"cc");
	return (int) res;
#endif
}

/*
 * pg_popcount_sse42
 *		Returns the number of 1-bits in buf
 */
pg_attribute_no_sanitize_alignment()
static uint64
pg_popcount_sse42(const char *buf, int bytes)
{
	uint64		popcnt = 0;
	const uint64 *words = (const uint64 *) buf;

	while (bytes >= 8)
	{
		popcnt += pg_popcount64_sse42(*words++);
		bytes -= 8;
	}

	buf = (const char *) words;

	/* Process any remaining bytes */
	while (bytes--)
		popcnt += pg_number_of_ones[(unsigned char) *buf++];

	return popcnt;
}

/*
 * pg_popcount_masked_sse42
 *		Returns the number of 1-bits in buf after applying the mask to each byte
 */
pg_attribute_no_sanitize_alignment()
static uint64
pg_popcount_masked_sse42(const char *buf, int bytes, uint8 mask)
{
	uint64		popcnt = 0;
	uint64		maskv = ~UINT64CONST(0) / 0xFF * mask;
	const uint64 *words = (const uint64 *) buf;

	while (bytes >= 8)
	{
		popcnt += pg_popcount64_sse42(*words++ & maskv);
		bytes -= 8;
	}

	buf = (const char *) words;

	/* Process any remaining bytes */
	while (bytes--)
		popcnt += pg_number_of_ones[(unsigned char) *buf++ & mask];

	return popcnt;
}

#else							/* HAVE_X86_64_POPCNTQ */

/* prevent linker complaints about empty module */
extern int	pg_popcount_x86_dummy_variable;
int			pg_popcount_x86_dummy_variable = 0;

#endif							/* ! HAVE_X86_64_POPCNTQ */