AVX2 - 使用 SIMD 加速最近邻查找

一、介绍

    SIMD（Single Instruction Multiple Data）：一条指令处理多个数据。

    AVX2：Intel/AMD 的 256 位 SIMD 指令集，一个 __m256 寄存器可以放 8 个 float 或 8 个 int32_t。

    使用cpu指令并行计算8个数。

二、AoS 与 SoA

1. 数据是 AoS（Array of Structures） 布局：

内存中：lat0, lon0, lat1, lon1, lat2, lon2, ...

2. AVX2 的加减乘除是按位置逐元素运算的，它希望数据是 SoA（Structure of Arrays）：

向量 A: [lat0, lat1, lat2, lat3, lat4, lat5, lat6, lat7]

向量 B: [lon0, lon1, lon2, lon3, lon4, lon5, lon6, lon7]

所以我们需要把交错的 lat/lon 去交错（de-interleave），拆成两个纯向量。这是整个优化中最烧脑的部分。

三、示例解析

    实时gps经纬度与指定gps数组最近距离求解：

int find_nearest_index_avx2(const GeoPoint* points, size_t count,

                            float t_lat, float t_lon, float* out_min_dist) {

    // 空数组边界处理

    if (count == 0) {

        if (out_min_dist) *out_min_dist = FLT_MAX;

        return -1;

    }

    // ========== 1. 预计算经纬度到米的换算系数 ==========

    // 纬度方向：1 度纬度 ≈ 111000 米

    const float LAT_TO_M = 111000.0f;

    // 经度方向：1 度经度 ≈ 111000 * cos(纬度) 米

    const float LON_TO_M = 111000.0f * cosf(t_lat * 3.141592653589793f / 180.0f);

    // 广播成 8 份，变成 AVX2 向量

    const __m256 v_lat_m = _mm256_set1_ps(LAT_TO_M);

    const __m256 v_lon_m = _mm256_set1_ps(LON_TO_M);

    // ========== 2. 目标点广播 ==========

    __m256 v_t_lat = _mm256_set1_ps(t_lat);

    __m256 v_t_lon = _mm256_set1_ps(t_lon);

    // ========== 3. 初始化最小值追踪器 ==========

    // v_min_dist：8 个通道各自记录“到目前为止见过的最小距离平方”

    __m256 v_min_dist = _mm256_set1_ps(FLT_MAX);

    // v_min_idx：8 个通道各自记录“最小距离对应的全局下标”

    __m256i v_min_idx = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);

    // v_base_idx：当前这批 8 个点的起始全局下标，初始为 0

    __m256i v_base_idx = _mm256_setzero_si256();

    // v_eight：常量 8 的向量，用于每轮结束后把基准下标 +8

    const __m256i v_eight = _mm256_set1_epi32(8);

    // deint：AoS→SoA 去交错索引表

    const __m256i deint = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);

    // ========== 4. 主循环：每次处理 8 个点 ==========

    size_t i = 0;

    for (; i + 7 < count; i += 8) {

        /*

         * points 内存布局：

         *   lat0,lon0, lat1,lon1, lat2,lon2, lat3,lon3,

         *   lat4,lon4, lat5,lon5, lat6,lon6, lat7,lon7, ...

         */

        // ---------- 4.1 加载 8 个点到 4 个 128 位寄存器 ----------

        __m128 lo1 = _mm_loadu_ps((const float*)&points[i]);        // lat0, lon0, lat1, lon1

        __m128 hi1 = _mm_loadu_ps((const float*)&points[i + 2]);    // lat2, lon2, lat3, lon3

        __m128 lo2 = _mm_loadu_ps((const float*)&points[i + 4]);    // lat4, lon4, lat5, lon5

        __m128 hi2 = _mm_loadu_ps((const float*)&points[i + 6]);    // lat6, lon6, lat7, lon7

        // ---------- 4.2 拼接成 2 个 256 位向量 ----------

        __m256 v_chunk1 = _mm256_insertf128_ps(_mm256_castps128_ps256(lo1), hi1, 1);

        // 结果：lat0, lon0, lat1, lon1 | lat2, lon2, lat3, lon3

        __m256 v_chunk2 = _mm256_insertf128_ps(_mm256_castps128_ps256(lo2), hi2, 1);

        // 结果：lat4, lon4, lat5, lon5 | lat6, lon6, lat7, lon7

        // ---------- 4.3 AoS → SoA 去交错（最烧脑的一步） ----------

        // _MM_SHUFFLE(z,y,x,w) 编码：结果 = [a[w], a[x], b[y], b[z]]

        __m256 t1 = _mm256_shuffle_ps(v_chunk1, v_chunk2, _MM_SHUFFLE(2, 0, 2, 0));

        // t1 = [lat0,lat1,lat4,lat5, lat2,lat3,lat6,lat7]

        __m256 t2 = _mm256_shuffle_ps(v_chunk1, v_chunk2, _MM_SHUFFLE(3, 1, 3, 1));

        // t2 = [lon0,lon1,lon4,lon5, lon2,lon3,lon6,lon7]

        // 用 permutevar8x32_ps 按 deint 索引重排，得到最终 SoA

        __m256 v_lats = _mm256_permutevar8x32_ps(t1, deint);

        // v_lats = [lat0,lat1,lat2,lat3,lat4,lat5,lat6,lat7]

        __m256 v_lons = _mm256_permutevar8x32_ps(t2, deint);

        // v_lons = [lon0,lon1,lon2,lon3,lon4,lon5,lon6,lon7]

        // ---------- 4.4 计算 8 个点的距离平方 ----------

        __m256 d_lat = _mm256_mul_ps(_mm256_sub_ps(v_lats, v_t_lat), v_lat_m);

        __m256 d_lon = _mm256_mul_ps(_mm256_sub_ps(v_lons, v_t_lon), v_lon_m);

        __m256 dist_sq = _mm256_add_ps(_mm256_mul_ps(d_lat, d_lat),

                                       _mm256_mul_ps(d_lon, d_lon));

        // ---------- 4.5 更新最小距离和对应索引（无分支） ----------

        // 比较：新距离是否比当前最小距离更小

        __m256 cmp_mask = _mm256_cmp_ps(v_min_dist, dist_sq, _CMP_GT_OQ);

        // cmp_mask 每个通道：真 = 0xFFFFFFFF，假 = 0x00000000

        // 更新最小距离：掩码为真 → 取新距离，否则取旧距离

        v_min_dist = _mm256_blendv_ps(v_min_dist, dist_sq, cmp_mask);

        // 更新索引：整数没有 blendv，用位运算实现同样的选择

        __m256i mask_i = _mm256_castps_si256(cmp_mask); // 按位重新解释为整数向量

        __m256i v_cur_idx = _mm256_add_epi32(v_base_idx,

                              _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));

        v_min_idx = _mm256_or_si256(

                        _mm256_and_si256(mask_i, v_cur_idx),      // 掩码为真 → 新下标

                        _mm256_andnot_si256(mask_i, v_min_idx));  // 掩码为假 → 旧下标

        // 基准下标 +8，准备下一轮

        v_base_idx = _mm256_add_epi32(v_base_idx, v_eight);

    }

    // ========== 5. 水平归约：从 8 个通道中找出最终答案 ==========

    ALIGN32 float min_dists[8];

    ALIGN32 int32_t min_idxs[8];

    _mm256_store_ps(min_dists, v_min_dist);

    _mm256_store_si256((__m256i*)min_idxs, v_min_idx);

    float final_min = min_dists[0];

    int final_idx = min_idxs[0];

    for (int k = 1; k < 8; k++) {

        if (min_dists[k] < final_min) {

            final_min = min_dists[k];

            final_idx = min_idxs[k];

        }

    }

    // ========== 6. 处理尾部不足 8 个的点 ==========

    for (; i < count; i++) {

        float dlat = (points[i].lat - t_lat) * LAT_TO_M;

        float dlon = (points[i].lon - t_lon) * LON_TO_M;

        float dist_sq = dlat * dlat + dlon * dlon;

        if (dist_sq < final_min) {

            final_min = dist_sq;

            final_idx = (int)i;

        }

    }

    // ========== 7. 输出结果 ==========

    if (out_min_dist) *out_min_dist = sqrtf(final_min);

    return final_idx;

}

四、函数详解

1. _mm256_sub_ps(a, b)

逐通道做 a - b，一次算 8 个 float 减法。

a = [10,20,30,40,50,60,70,80]

b = [ 1, 2, 3, 4, 5, 6, 7, 8]

c = _mm256_sub_ps(a, b); // [9,18,27,36,45,54,63,72]

2. _mm256_add_epi32(a, b)

逐通道做 32 位整数加法，一次算 8 个 int32_t。

a = [1,2,3,4,5,6,7,8]

b = [10,20,30,40,50,60,70,80]

c = _mm256_add_epi32(a, b); // [11,22,33,44,55,66,77,88]

3. _mm256_setr_epi32(0,1,2,3,4,5,6,7)

创建一个整数向量，参数从低到高填入。setr 的 r 表示 reversed，参数顺序符合直觉（低→高）。对应的 _mm256_set_epi32 则是从高到低填，容易搞混。

4. _mm256_castps_si256(x)

按位重新解释，把 __m256（8 个 float）看成 __m256i（8 个 int32），二进制内容不变，不是数值转换。零开销。

5. 无分支索引更新

v_min_idx = (mask & v_cur_idx) | (~mask & v_min_idx);

mask 全 1 → 取 v_cur_idx（更新）

mask 全 0 → 取 v_min_idx（保留）

等价于 if (mask) cur = new; else cur = old;，但一次处理 8 个通道，且没有分支。

五、示例总结

        分为两部分

        1. 原始数据内存结构

            lat0,lon0,lat1,lon1,lat2,lon2,lat3,lon3,lat4,lon4,lat5,lon5,lat6,lon6,lat7,lon7,

        转换两个通道

        0通道

            lat0,lat1,lat2,lat3, lat4,lat5,lat6,lat7

        1通道

            lon0,lon1,lon2,lon3, lon4,lon5,lon6,lon7

        2. 数值比较

            保留通道每个数据最小值，保留通道每个数据的最小值的索引

            循环与对应位置的历史最小值对比，同时保留最小值索引

        3. 常规方法取出最小值和索引