我在一台机器上拥有一个在热路径中运行、对性能极为敏感的半径搜索函数，运行在一个 Intel Xeon Ice Lake 8352Y（支持AVX-512，32核，2.2 GHz）上，使用GCC 14.3.0编译。

该函数对以SoA形式组织的三维坐标数组计算平方欧氏距离，并将半径内点的索引和距离存储起来。

unsigned int radius_search_avx512(const float *restrict xs, const float *restrict ys, const float *restrict zs,
                 unsigned int window, unsigned int search_start_index, float x, float y, float z,
                 float radius, unsigned int *restrict indices, float *restrict distances)
{
    unsigned int elements_count = 0;

    // Create a vector with the position of the point to compare (x0, y0, z0)
    const __m512 x0_vector = _mm512_set1_ps(x);
    const __m512 y0_vector = _mm512_set1_ps(y);
    const __m512 z0_vector = _mm512_set1_ps(z);

    // Create a vector with the radius
    __m512 radius_vector = _mm512_set1_ps(radius);

    __m512i increment = _mm512_set1_epi32(16);

    // Create base index (where search starts)
    __m512i base_idx = _mm512_set_epi32(
        search_start_index + 15, search_start_index + 14, search_start_index + 13, search_start_index + 12,
        search_start_index + 11, search_start_index + 10, search_start_index + 9, search_start_index + 8,
        search_start_index + 7, search_start_index + 6, search_start_index + 5, search_start_index + 4,
        search_start_index + 3, search_start_index + 2, search_start_index + 1, search_start_index + 0);

    for (unsigned int i = 0; i + 15 < window; i += 16) {
        // Load + differences
        __m512 x_vector_result = _mm512_sub_ps(_mm512_load_ps(xs + i), x0_vector);
        __m512 y_vector_result = _mm512_sub_ps(_mm512_load_ps(ys + i), y0_vector);
        __m512 z_vector_result = _mm512_sub_ps(_mm512_load_ps(zs + i), z0_vector);

        // Squared Euclidean distance
        __m512 result_vec = _mm512_mul_ps(x_vector_result, x_vector_result);
        result_vec = _mm512_fmadd_ps(y_vector_result, y_vector_result, result_vec);
        result_vec = _mm512_fmadd_ps(z_vector_result, z_vector_result, result_vec);

        // Compare distance <= radius
        __mmask16 mask = _mm512_cmp_ps_mask(result_vec, radius_vector, _CMP_LE_OQ);

        if (mask != 0) {
            // Save matching results
            _mm512_mask_compressstoreu_epi32(indices + elements_count, mask, base_idx);
            _mm512_mask_compressstoreu_ps(distances + elements_count, mask, result_vec);

            // Count valid elements
            elements_count += __builtin_popcount(mask);
        }

        // Increment indices
        base_idx = _mm512_add_epi32(base_idx, increment);
    }

    return elements_count;
}

背景 / 观察

• 所有输入数组都按64字节对齐。
• 使用启用AVX-512的优化标志进行编译。
• 该函数处于一个非常热的路径中。
• 通过分析，主要瓶颈似乎是内存加载，因为带宽受限。

• 我已经尝试过：

• 取消对 base_idx 的依赖

• 重新排序操作
• 避免与其他操作混合使用compressstoreu

• 预取

• 由于压缩导致的依赖，elements_count 看起来不可避免。

• 条件 if (mask != 0) 通常基本为假
• 数组非常大，包含数百万个浮点数（从10到 4亿不等）。不过窗口很小（在某些情况下约1000或 10000）。
• 该函数在for循环中对数组中的每一个点按顺序调用，并使用OpenMP做并行。

还有没有进一步的优化办法？

解决方案

该函数受内存带宽限制。

每次迭代需要3 次加载是瓶颈，因此为了减少加载，可以进行提前剪枝，检查 (x² > radius) 和 (x²+y² > radius)，从而避免一些加载。由于条件 if (mask != 0) 通常基本为假，这种剪枝检查通常是有效的。

评论中提到的其他优化也有效1。

最终代码在不同数据集上的平均速度提升约为x3：

index_t tfg_radius_search_avx512(const float *restrict xs, const float *restrict ys, const float *restrict zs,
                 unsigned int window, unsigned int search_start_index, float x, float y, float z,
                 float radius, unsigned int *restrict indices, float *restrict distances)
{
    unsigned int elements_count = 0;

    // Create a vector with the position of the point to compare (x0, y0, z0)
    const __m512 x0_vector = _mm512_set1_ps(x);
    const __m512 y0_vector = _mm512_set1_ps(y);
    const __m512 z0_vector = _mm512_set1_ps(z);

    // Create a vector with the radius
    const __m512 radius_vector = _mm512_set1_ps(radius);

    const __m512i increment = _mm512_set1_epi32(16);
    const __m512i increment_32 = _mm512_set1_epi32(32);

    // Create base index (where search starts)
    const __m512i idx_offsets = _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0);
    __m512i base_idx = _mm512_add_epi32(_mm512_set1_epi32(search_start_index), idx_offsets);

    for (unsigned int i = 0; i + 31 < window; i += 32) {
        // X²
        __m512 x_vector_result = _mm512_sub_ps(x0_vector, _mm512_load_ps(xs + i));
        __m512 x2_vector_result = _mm512_sub_ps(x0_vector, _mm512_load_ps(xs + i + 16));
        __m512 result_vec = _mm512_mul_ps(x_vector_result, x_vector_result);
        __m512 result_vec2 = _mm512_mul_ps(x2_vector_result, x2_vector_result);

        if (_mm512_cmp_ps_mask(result_vec, radius_vector, _CMP_LE_OQ) |
            _mm512_cmp_ps_mask(result_vec2, radius_vector, _CMP_LE_OQ)) {
                // X² + Y²
            __m512 y_vector_result = _mm512_sub_ps(y0_vector, _mm512_load_ps(ys + i));
            __m512 y2_vector_result = _mm512_sub_ps(y0_vector, _mm512_load_ps(ys + i + 16));
            result_vec = _mm512_fmadd_ps(y_vector_result, y_vector_result, result_vec);
            result_vec2 = _mm512_fmadd_ps(y2_vector_result, y2_vector_result, result_vec2);

            if (_mm512_cmp_ps_mask(result_vec, radius_vector, _CMP_LE_OQ) |
                _mm512_cmp_ps_mask(result_vec2, radius_vector, _CMP_LE_OQ)) {
                    // X² + Y² + Z²
                __m512 z_vector_result = _mm512_sub_ps(z0_vector, _mm512_load_ps(zs + i));
                __m512 z2_vector_result = _mm512_sub_ps(z0_vector, _mm512_load_ps(zs + i + 16));
                result_vec = _mm512_fmadd_ps(z_vector_result, z_vector_result, result_vec);
                result_vec2 = _mm512_fmadd_ps(z2_vector_result, z2_vector_result, result_vec2);

                // Compare distance <= radius
                __mmask16 mask = _mm512_cmp_ps_mask(result_vec, radius_vector, _CMP_LE_OQ);
                __mmask16 mask2 = _mm512_cmp_ps_mask(result_vec2, radius_vector, _CMP_LE_OQ);

                if (mask) {
                    // Save matching results
                    _mm512_mask_compressstoreu_epi32(indices + elements_count, mask, base_idx);
                    _mm512_mask_compressstoreu_ps(distances + elements_count, mask, result_vec);

                    // Count valid elements
                    elements_count += __builtin_popcount(mask);
                }
                if (mask2) {
                    // Save matching results
                    _mm512_mask_compressstoreu_epi32(indices + elements_count, mask2,
                                     _mm512_add_epi32(base_idx, increment));
                    _mm512_mask_compressstoreu_ps(distances + elements_count, mask2, result_vec2);

                    // Count valid elements
                    elements_count += __builtin_popcount(mask2);
                }
            }
        }
        // Increment indices
        base_idx = _mm512_add_epi32(base_idx, increment_32);
    }

    return elements_count;
}

在实践中，x2的展开似乎效果更好。

脚注1：来自 @PeterCordes与 @JérômeRichard的微优化

• _mm512_sub_pd(vec, load()); 代替另一操作数顺序，允许一个内存源操作数折叠到 vsubpd 指令中，而不需要单独的 vmovapd 加载。翻转符号也没问题，因为我们只对结果进行平方。这仅有助于前端吞吐量，并让ROB（重排序缓冲区）的空间使用更少，从而让乱序执行看得更远，而不会影响后端执行吞吐资源，如执行端口。

• _mm512_set_epi32( search_start_index + 15, ...) 在GCC上表现很差。只有Clang能将其优化为 vpbroadcastd + vpaddd，并带有一个常量。GCC会对16个整数逐个执行标量运算以及vpinsrd等。https://godbolt.org/z/MbMrj6fr9

• 使用（或让编译器使用）kortest 来一次性查看两个掩码中是否有任意位被置位，因此用 if(mask1 | mask2) 代替仅仅 if(mask)。根据分支预测细节和数据分布（如块内聚簇），你可以对两个集合无条件执行 compressstoreu，或增加分支。

• compressstoreu 在英特尔CPU上表现良好，和 vcompressps zmm{k}{z},zmm 的uop数量相同，并且需要一个单独且未屏蔽的 vmovups ZMM存储。在AMD CPU上，单独的存储更优，因为为 vcompressps 的内存目标在Zen 4/5上会花费8 个uops而不是2。https://uops.info/ 。( vmovups m512, zmm 在Zen 4/5上的存储成本为2 个uops；存储队列条目只有32字节。所以对向量寄存器进行压缩再存储，总共需要4 个uops，并且可以与前面的存储重叠。)

• 对齐到64字节的加载已经是内存带宽的最佳情况，硬件预取应该能工作，因此没有必要做软件预取。但如果在数据窗口已经热在缓存中时再这样做，将有很大帮助。因此，基本思路是在调用此函数及其它在小窗口内运行的函数的外层循环中进行缓存块化。（特别是在同一核心上的L2缓存，但英特尔的共享L3让任何核心都能受益，尤其是在单插槽系统中。）

*（对于条件大多数情况为真的其他用例）—— 如果你存储大量数据（表面上看起来确实如此），可以使用流式存储（也称为非临时存储）。你可以把带掩码的项写入一个容量较小、能容纳在L1缓存中的临时缓冲区，当此缓冲区满时再将数据流式写入主存。这可以节省一些内存带宽，因为常规存储会触发缓存行写分配，由于不必要的读取而浪费带宽。