我有一个性能关键的半径搜索函数，运行在Intel Xeon Ice Lake 8352Y（支持AVX-512、32核、2.2 GHz）平台上，使用gcc 14.3.0编译，处于热执行路径中。该函数计算SoA格式3D坐标数组的平方欧氏距离，并存储给定半径内点的索引与距离。

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;

    // 创建待比较点的向量(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);

    // 创建半径向量
    __m512 radius_vector = _mm512_set1_ps(radius);

    __m512i increment = _mm512_set1_epi32(16);

    // 创建起始索引向量
    __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) {
        // 加载坐标并计算差值
        __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);

        // 计算平方欧氏距离
        __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);

        // 比较距离是否小于等于半径
        __mmask16 mask = _mm512_cmp_ps_mask(result_vec, radius_vector, _CMP_LE_OQ);

        if (mask != 0) {
            // 保存匹配结果
            _mm512_mask_compressstoreu_epi32(indices + elements_count, mask, base_idx);
            _mm512_mask_compressstoreu_ps(distances + elements_count, mask, result_vec);

            // 统计有效元素数量
            elements_count += __builtin_popcount(mask);
        }

        // 更新索引向量
        base_idx = _mm512_add_epi32(base_idx, increment);
    }

    return elements_count;
}

上下文与观察结果 #

• 所有输入数组均为64字节对齐。
• 编译时已启用AVX-512优化选项。
• 该函数处于极热执行路径中。
• 性能分析显示主要瓶颈为内存加载，这令我感到困惑。
• 已尝试以下优化：
  • 移除对base_idx的依赖
  • 重排操作顺序
  • 替换compressstoreu操作
  • 预取数据
• 由于需要压缩结果，elements_count的依赖似乎无法避免。

进一步优化方案 #

1. 数据布局切换为AoSoA

当前纯SoA布局（x、y、z分存三个数组）会导致内存访问分散，每次需要加载3个独立的512位向量。改为AoSoA（数组结构的数组）可以将一组点的x/y/z数据打包存储，减少内存请求次数，提升缓存命中率：

struct alignas(64) PointBlock {
    float x[16];
    float y[16];
    float z[16];
};

加载时一次性读取整个PointBlock，再分别提取x/y/z向量，降低内存总线压力。

2. 流水线重叠：计算与内存加载并行

将内存加载与计算逻辑重排，利用计算时间掩盖内存加载延迟：

// 预加载第一组数据
__m512 x_vec = _mm512_load_ps(xs);
__m512 y_vec = _mm512_load_ps(ys);
__m512 z_vec = _mm512_load_ps(zs);

for (unsigned int i = 16; i <= window; i += 16) {
    // 计算当前组距离
    __m512 x_diff = _mm512_sub_ps(x_vec, x0_vector);
    __m512 y_diff = _mm512_sub_ps(y_vec, y0_vector);
    __m512 z_diff = _mm512_sub_ps(z_vec, z0_vector);

    __m512 dist_sq = _mm512_mul_ps(x_diff, x_diff);
    dist_sq = _mm512_fmadd_ps(y_diff, y_diff, dist_sq);
    dist_sq = _mm512_fmadd_ps(z_diff, z_diff, dist_sq);

    // 预加载下一组数据
    __m512 next_x = _mm512_load_ps(xs + i);
    __m512 next_y = _mm512_load_ps(ys + i);
    __m512 next_z = _mm512_load_ps(zs + i);

    // 处理当前组匹配结果
    __mmask16 mask = _mm512_cmp_ps_mask(dist_sq, radius_vector, _CMP_LE_OQ);
    if (mask != 0) {
        _mm512_mask_compressstoreu_epi32(indices + elements_count, mask, base_idx);
        _mm512_mask_compressstoreu_ps(distances + elements_count, mask, dist_sq);
        elements_count += __builtin_popcount(mask);
    }

    // 更新当前向量与索引
    x_vec = next_x;
    y_vec = next_y;
    z_vec = next_z;
    base_idx = _mm512_add_epi32(base_idx, increment);
}

3. 预计算半径平方

将半径提前平方，避免比较时的精度问题，同时让逻辑更准确：

float radius_sq = radius * radius;
__m512 radius_vector = _mm512_set1_ps(radius_sq);

4. 精细化缓存预取

调整预取时机与层级，避免缓存污染，比如在循环中提前预取下下组数据到L1缓存：

_mm_prefetch((const char*)(xs + i + 32), _MM_HINT_T0);
_mm_prefetch((const char*)(ys + i + 32), _MM_HINT_T0);
_mm_prefetch((const char*)(zs + i + 32), _MM_HINT_T0);

5. 批量处理多搜索点（若场景允许）

如果函数需频繁处理不同搜索点，可改为一次处理多个搜索点，利用AVX-512向量宽度提升数据复用率，减少重复内存加载。

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