如何基于AVX指令集实现单精度数组的12连续元素均值计算及结果6次复制功能？

Great question! Using AVX to optimize this kind of block-wise averaging and replication is a smart move for throughput, especially with large arrays. Let's break down how to implement this efficiently with Intel intrinsics (the easiest way to work with AVX without raw assembly):

Core Approach #

Your task splits into two key steps per 12-element input block:

1. Compute the mean: Sum 12 single-precision floats, then divide by 12.
2. Replicate the mean: Write the same mean value 6 times to the output array.

AVX 256-bit registers can hold 8 floats at once, so we'll leverage that to minimize memory access overhead.

Implementation Code #

Here's a complete, optimized implementation using AVX/AVX2 intrinsics:

#include <immintrin.h>
#include <string.h>

// input: array of length 12 * num_blocks (single-precision floats)
// output: array of length 6 * num_blocks (single-precision floats)
void process_avx(const float* input, float* output, size_t num_blocks) {
    // Precompute reciprocal of 12 for faster multiplication (vs division)
    const float inv_12 = 1.0f / 12.0f;

    for (size_t i = 0; i < num_blocks; ++i) {
        const float* curr_input = input + (i * 12);
        float* curr_output = output + (i * 6);

        // Step 1: Sum the 12 input elements
        // Load first 8 elements into a 256-bit AVX register
        __m256 vec_8 = _mm256_loadu_ps(curr_input);
        // Sum all 8 elements (use AVX2 reduce if available for cleaner code)
        float sum_8;
#ifdef __AVX2__
        sum_8 = _mm256_reduce_add_ps(vec_8);
#else
        // Fallback for AVX (no reduce instruction)
        __m128 low_128 = _mm256_extractf128_ps(vec_8, 0);
        __m128 high_128 = _mm256_extractf128_ps(vec_8, 1);
        __m128 sum_128 = _mm_add_ps(low_128, high_128);
        sum_128 = _mm_hadd_ps(sum_128, sum_128);
        sum_128 = _mm_hadd_ps(sum_128, sum_128);
        sum_8 = _mm_cvtss_f32(sum_128);
#endif

        // Load remaining 4 elements into a 128-bit SSE register
        __m128 vec_4 = _mm_loadu_ps(curr_input + 8);
        // Sum the 4 elements
        float sum_4;
#ifdef __AVX2__
        sum_4 = _mm_reduce_add_ps(vec_4);
#else
        __m128 sum_4_128 = _mm_hadd_ps(vec_4, vec_4);
        sum_4_128 = _mm_hadd_ps(sum_4_128, sum_4_128);
        sum_4 = _mm_cvtss_f32(sum_4_128);
#endif

        // Calculate mean
        float total_sum = sum_8 + sum_4;
        float mean = total_sum * inv_12;

        // Step 2: Write mean 6 times to output
        // Broadcast mean to a 256-bit register (8 copies of the mean)
        __m256 mean_broadcast = _mm256_set1_ps(mean);
        // Write first 4 elements to output
        _mm_storeu_ps(curr_output, _mm256_castps256_ps128(mean_broadcast));
        // Write remaining 2 elements (use memcpy for simplicity, or intrinsics)
        memcpy(curr_output + 4, &mean_broadcast, 2 * sizeof(float));
    }
}

Key Optimizations Explained #

• Reciprocal Multiplication: We use inv_12 = 1/12 instead of dividing by 12 directly—floating-point multiplication is significantly faster than division.
• AVX Memory Operations: _mm256_loadu_ps/_mm_storeu_ps handle unaligned memory; if your input/output arrays are 32-byte aligned (use aligned_alloc or posix_memalign), switch to _mm256_load_ps/_mm256_store_ps for even better performance.
• AVX2 Reduce Instructions: If AVX2 is available (most modern x86 CPUs), _mm256_reduce_add_ps and _mm_reduce_add_ps simplify and speed up horizontal summation compared to manual hadd operations.
• Broadcast & Batch Writes: Broadcasting the mean to a 256-bit register lets us write 4 elements at once, cutting down on memory write overhead.

Further Optimizations #

• Loop Unrolling: Unroll the loop to process 2 or 4 blocks at a time (e.g., 24 or 48 input elements) to hide memory access latency and reduce loop control overhead.
• Vectorized Mean Calculation: For even more throughput, you could process multiple 12-element blocks in parallel using additional AVX registers, but this requires careful register allocation.

内容的提问来源于stack exchange，提问作者ACE