Quelle  pg_popcount_avx512.c

/*-------------------------------------------------------------------------
 *
 * pg_popcount_avx512.c
 *   Holds the AVX-512 pg_popcount() implementation.
 *
 * Copyright (c) 2024-2025, PostgreSQL Global Development Group
 *
 * IDENTIFICATION
 *   src/port/pg_popcount_avx512.c
 *
 *-------------------------------------------------------------------------
 */
#include "c.h"

#ifdef USE_AVX512_POPCNT_WITH_RUNTIME_CHECK

#if defined(HAVE__GET_CPUID) || defined(HAVE__GET_CPUID_COUNT)
#include <cpuid.h>
#endif

#include <immintrin.h>

#if defined(HAVE__CPUID) || defined(HAVE__CPUIDEX)
#include <intrin.h>
#endif

#include "port/pg_bitutils.h"

/*
 * It's probably unlikely that TRY_POPCNT_X86_64 won't be set if we are able to
 * use AVX-512 intrinsics, but we check it anyway to be sure.  We piggy-back on
 * the function pointers that are only used when TRY_POPCNT_X86_64 is set.
 */
#ifdef TRY_POPCNT_X86_64

/*
 * Does CPUID say there's support for XSAVE instructions?
 */
static inline bool
xsave_available(void)
{
 unsigned int exx[4] = {0, 0, 0, 0};

#if defined(HAVE__GET_CPUID)
 __get_cpuid(1, &exx[0], &exx[1], &exx[le='color: green'>2], &exx[3]);
#elif defined(HAVE__CPUID)
 __cpuid(exx, 1);
#else
#error cpuid instruction not available
#endif
 return (exx[2] & (1 << 27)) != 0; /* osxsave */
}

/*
 * Does XGETBV say the ZMM registers are enabled?
 *
 * NB: Caller is responsible for verifying that xsave_available() returns true
 * before calling this.
 */
#ifdef HAVE_XSAVE_INTRINSICS
pg_attribute_target("xsave")
#endif
static inline bool
zmm_regs_available(void)
{
#ifdef HAVE_XSAVE_INTRINSICS
 return (_xgetbv(0) & 0xe6) == 0xe6;
#else
 return false;
#endif
}

/*
 * Does CPUID say there's support for AVX-512 popcount and byte-and-word
 * instructions?
 */
static inline bool
avx512_popcnt_available(void)
{
 unsigned int exx[4] = {0, 0, 0, 0};

#if defined(HAVE__GET_CPUID_COUNT)
 __get_cpuid_count(7, 0, &exx[0], &exx[1], &exx[<span style='color: green'>2], &exx[3]);
#elif defined(HAVE__CPUIDEX)
 __cpuidex(exx, 7, 0);
#else
#error cpuid instruction not available
#endif
 return (exx[2] & (1 << 14)) != 0 && /* avx512-vpopcntdq */
  (exx[1] & (1 << 30)) != 0; /* avx512-bw */
}

/*
 * Returns true if the CPU supports the instructions required for the AVX-512
 * pg_popcount() implementation.
 */
bool
pg_popcount_avx512_available(void)
{
 return xsave_available() &&
  zmm_regs_available() &&
  avx512_popcnt_available();
}

/*
 * pg_popcount_avx512
 *  Returns the number of 1-bits in buf
 */
pg_attribute_target("avx512vpopcntdq,avx512bw")
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")
uint64
pg_popcount_masked_avx512(const char *buf, int bytes, bits8 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       /* TRY_POPCNT_X86_64 */

#else       /* USE_AVX512_POPCNT_WITH_RUNTIME_CHECK */

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

#endif       /* ! USE_AVX512_POPCNT_WITH_RUNTIME_CHECK */