Repository: valkey-io/valkey
File: src/hyperloglog.c (C)

Code before:
        y1 = _mm256_or_si256(a1, a2);
        y2 = _mm256_or_si256(a3, a4);
        y = _mm256_or_si256(y1, y2);

        __m256i z = _mm256_loadu_si256((__m256i *)t);

        z = _mm256_max_epu8(z, y);

        _mm256_storeu_si256((__m256i *)t, z);

        r += 24;
        t += 32;
    }

    /* Merge the last 24 registers normally
     * as the AVX2 algorithm needs 4 padding bytes at the end */
    for (int i = HLL_REGISTERS - 24; i < HLL_REGISTERS; i++) {
        HLL_DENSE_GET_REGISTER(val, reg_dense, i);
        if (val > reg_raw[i]) {
            reg_raw[i] = val;
        }
    }
}
#endif

#if defined(__aarch64__)
/*
 * hllMergeDenseNEON is an ARM optimized version of hllMergeDense using NEON
 *
 * This function merges HyperLogLog (HLL) dense registers using ARM NEON SIMD instructions.
 * It extracts 6 bits registers from a dense format, and stores them in raw format
 *
 * Parameters:
 * - reg_raw: Pointer to the raw register array
 * - reg_dense: Pointer to the dense register array
 */
void hllMergeDenseNEON(uint8_t *reg_raw, const uint8_t *reg_dense) {
    uint8_t *dense_ptr = (uint8_t *)reg_dense;
    uint8_t *raw_ptr = (uint8_t *)reg_raw;

    uint8x16_t idx = {0, 1, 2, 0xFF,
                      3, 4, 5, 0xFF,
                      6, 7, 8, 0xFF,
                      9, 10, 11, 0xFF};

    // Bit masks for extracting specific bit ranges
    uint8x16_t mask1 = vreinterpretq_u8_u32(vdupq_n_u32(0x0000003f)); // Bits 0-5
    uint8x16_t mask2 = vreinterpretq_u8_u32(vdupq_n_u32(0x00000fc0)); // Bits 6-11
    uint8x16_t mask3 = vreinterpretq_u8_u32(vdupq_n_u32(0x0003f000)); // Bits 12-17
    uint8x16_t mask4 = vreinterpretq_u8_u32(vdupq_n_u32(0x00fc0000)); // Bits 18-23


Review comment (suggestion):
Other conditional blocks in this file use

    #if defined(__aarch64__) && HAVE_ARM_NEON

Should it be the same check in all places?

To test this properly, we need a 32-bit ARM machine with NEON. I don't think we have one.... Maybe we could cross-compile, just to see that it compiles.

Code after:
        y1 = _mm256_or_si256(a1, a2);
        y2 = _mm256_or_si256(a3, a4);
        y = _mm256_or_si256(y1, y2);

        __m256i z = _mm256_loadu_si256((__m256i *)t);

        z = _mm256_max_epu8(z, y);

        _mm256_storeu_si256((__m256i *)t, z);

        r += 24;
        t += 32;
    }

    /* Merge the last 24 registers normally
     * as the AVX2 algorithm needs 4 padding bytes at the end */
    for (int i = HLL_REGISTERS - 24; i < HLL_REGISTERS; i++) {
        HLL_DENSE_GET_REGISTER(val, reg_dense, i);
        if (val > reg_raw[i]) {
            reg_raw[i] = val;
        }
    }
}
#endif

#if HAVE_ARM_NEON
/*
 * hllMergeDenseNEON is an ARM optimized version of hllMergeDense using NEON
 *
 * This function merges HyperLogLog (HLL) dense registers using ARM NEON SIMD instructions.
 * It extracts 6 bits registers from a dense format, and stores them in raw format
 *
 * Parameters:
 * - reg_raw: Pointer to the raw register array
 * - reg_dense: Pointer to the dense register array
 */
void hllMergeDenseNEON(uint8_t *reg_raw, const uint8_t *reg_dense) {
    uint8_t *dense_ptr = (uint8_t *)reg_dense;
    uint8_t *raw_ptr = (uint8_t *)reg_raw;

    uint8x16_t idx = {0, 1, 2, 0xFF,
                      3, 4, 5, 0xFF,
                      6, 7, 8, 0xFF,
                      9, 10, 11, 0xFF};

    // Bit masks for extracting specific bit ranges
    uint8x16_t mask1 = vreinterpretq_u8_u32(vdupq_n_u32(0x0000003f)); // Bits 0-5
    uint8x16_t mask2 = vreinterpretq_u8_u32(vdupq_n_u32(0x00000fc0)); // Bits 6-11
    uint8x16_t mask3 = vreinterpretq_u8_u32(vdupq_n_u32(0x0003f000)); // Bits 12-17
    uint8x16_t mask4 = vreinterpretq_u8_u32(vdupq_n_u32(0x00fc0000)); // Bits 18-23