OpenMP+AVX版本基数排序结果异常，但仅OpenMP版本正常，求排查原因

我正在实现一个针对64位无符号整数（key_t 为 uint_fast64_t）的基数排序，这些整数是编码后的点数据。目前有两个实现版本：

• 仅使用OpenMP的版本 ✅ 排序结果完全正确
• 结合OpenMP+AVX2的版本 ❌ 输出排序结果错误

两个版本均通过OpenMP对输入数组分块并行处理，问题出现在我使用AVX指令向量化直方图计数和元素重分配阶段的逻辑中。

仅OpenMP版本代码 #

template <typename Point_t>
std::vector<key_t> mySort_v1(std::vector<Point_t> &points, std::optional<std::vector<PointMetadata>> &meta_opt, const Box &bbox) const {
    size_t n = points.size();
    constexpr int BITS_PER_PASS = 8;
    constexpr int NUM_BUCKETS = 1 << BITS_PER_PASS;
    constexpr size_t BUCKET_MASK = NUM_BUCKETS - 1;
    constexpr int NUM_PASSES = sizeof(key_t) * 8 / BITS_PER_PASS;

    // Encoding
    std::vector<key_t> keys(n);
    #pragma omp parallel for schedule(static)
    for (size_t i = 0; i < n; ++i) {
        keys[i] = encodeFromPoint(points[i], bbox);
    }

    std::vector<key_t> buffer(n);
    std::vector<PointMetadata> metadata_buffer;
    std::vector<Point_t> bufferDecoded(n);
    if (meta_opt) metadata_buffer.resize(n);

    for (int pass = 0; pass < NUM_PASSES; pass++) {
        int shift = pass * BITS_PER_PASS;
        const int nThreads = omp_get_max_threads();
        std::vector<std::vector<size_t>> localHist(nThreads, std::vector<size_t>(NUM_BUCKETS, 0));

        // Step 1: Histogram counting
        #pragma omp parallel
        {
            auto &hist = localHist[omp_get_thread_num()];
            #pragma omp for nowait schedule(static)
            for (size_t i = 0; i < n; ++i) {
                size_t bucket = (keys[i] >> shift) & BUCKET_MASK;
                hist[bucket]++;
            }
        }

        // Step 2: Compute bucket offsets from histograms
        size_t offset = 0;
        for (int b = 0; b < NUM_BUCKETS; b++) {
            for (int t = 0; t < nThreads; t++) {
                size_t val = localHist[t][b];
                localHist[t][b] = offset;
                offset += val;
            }
        }

        // Step 3: Scatter elements to buffer using per-thread offsets
        #pragma omp parallel
        {
            auto &localOffset = localHist[omp_get_thread_num()];
            #pragma omp for schedule(static)
            for (size_t i = 0; i < n; i++) {
                size_t bucket = (keys[i] >> shift) & BUCKET_MASK;
                size_t pos = localOffset[bucket]++;
                buffer[pos] = keys[i];
                bufferDecoded[pos] = points[i];
                if (meta_opt) metadata_buffer[pos] = (*meta_opt)[i];
            }
        }

        std::swap(points, bufferDecoded);
        std::swap(keys, buffer);
        if (meta_opt) std::swap(*meta_opt, metadata_buffer);
    }
    return keys;
}

OpenMP+AVX版本代码（存在问题） #

template <typename Point_t>
std::vector<key_t> BasicOptimizedAVX(std::vector<Point_t> &points, std::optional<std::vector<PointMetadata>> &meta_opt, const Box &bbox) const {
    size_t n = points.size();
    constexpr int BITS_PER_PASS = 8;
    constexpr int NUM_BUCKETS = 1 << BITS_PER_PASS;
    constexpr size_t BUCKET_MASK = NUM_BUCKETS - 1;
    constexpr int NUM_PASSES = sizeof(key_t) * 8 / BITS_PER_PASS;

    // Encoding
    std::vector<key_t> keys(n);
    #pragma omp parallel for schedule(static)
    for (size_t i = 0; i < n; ++i) {
        keys[i] = encodeFromPoint(points[i], bbox);
    }

    std::vector<key_t> buffer(n);
    std::vector<PointMetadata> metadata_buffer;
    std::vector<Point_t> bufferDecoded(n);
    if (meta_opt) metadata_buffer.resize(n);

    // 重复mask_value 4次到256位向量
    __m256i mask_vec = _mm256_set1_epi64x((int64_t)BUCKET_MASK);

    for (int pass = 0; pass < NUM_PASSES; pass++) {
        int shift = pass * BITS_PER_PASS;
        const int nThreads = omp_get_max_threads();
        std::vector<std::vector<size_t>> localHist(nThreads, std::vector<size_t>(NUM_BUCKETS, 0));

        // Step 1: Histogram counting
        #pragma omp parallel
        {
            auto &hist = localHist[omp_get_thread_num()];
            #pragma omp for schedule(static)
            for (size_t i = 0; i <= n - 4; i += 4) {
                // 加载数据到AVX寄存器
                __m256i vec = _mm256_loadu_si256((__m256i*)&keys[i]);
                // 右移操作
                __m256i shifted_right = _mm256_srli_epi64(vec, shift);
                // 按位与操作
                __m256i result_and = _mm256_and_si256(shifted_right, mask_vec);
                // 存储结果
                alignas(32) uint64_t result[4];
                _mm256_storeu_si256((__m256i *)result, result_and);
                // 更新直方图
                hist[result[0]]++;
                hist[result[1]]++;
                hist[result[2]]++;
                hist[result[3]]++;
            }
        }

        // 处理剩余不足4个的元素
        for (size_t i = (n / 4) * 4; i < n; ++i) {
            size_t bucket = (keys[i] >> shift) & BUCKET_MASK;
            localHist[0][bucket]++;
        }

        // Step 2: Compute bucket offsets from histograms
        size_t offset = 0;
        for (int b = 0; b < NUM_BUCKETS; b++) {
            for (int t = 0; t < nThreads; t++) {
                size_t val = localHist[t][b];
                localHist[t][b] = offset;
                offset += val;
            }
        }

        // Step 3: Scatter elements to buffer using per-thread offsets
        #pragma omp parallel
        {
            auto &localOffset = localHist[omp_get_thread_num()];
            #pragma omp for schedule(static)
            for (size_t i = 0; i <= n - 4; i += 4) {
                __m256i vec = _mm256_loadu_si256((__m256i*)&keys[i]);
                // Shift >>
                __m256i shifted_right = _mm256_srli_epi64(vec, shift);
                __m256i result_and = _mm256_and_si256(shifted_right, mask_vec);
                alignas(32) uint64_t buckets[4];
                _mm256_storeu_si256((__m256i *)buckets, result_and);
                // 【代码不完整，此处逻辑缺失】
            }
            // 【剩余元素处理逻辑缺失】
        }

        std::swap(points, bufferDecoded);
        std::swap(keys, buffer);
        if (meta_opt) std::swap(*meta_opt, metadata_buffer);
    }
    return keys;
}

问题根源分析 #

1. 剩余元素的直方图统计逻辑错误 #

AVX版本中，剩余元素的处理循环位于OpenMP parallel区域之外，且直接将计数加到localHist[0]（线程0的直方图）：

// 处理剩余不足4个的元素
for (size_t i = (n / 4) * 4; i < n; ++i) {
    size_t bucket = (keys[i] >> shift) & BUCKET_MASK;
    localHist[0][bucket]++;
}

这会导致两个致命问题：

• 剩余元素的计数仅被统计到线程0的直方图中，其他线程的直方图未包含这些元素的计数，后续的偏移计算（Step2）会出现偏差，最终导致元素重分配时的位置错误。
• 剩余元素未被并行处理，不仅效率低下，还破坏了直方图统计的正确性。

2. 元素重分配阶段的AVX代码不完整且缺失剩余元素处理 #

从提供的代码来看，重分配阶段（Step3）的AVX逻辑仅完成了bucket计算，缺失了：

• 根据bucket获取对应线程的偏移量并自增
• 将元素（keys、points、metadata）存储到buffer的对应位置
• 剩余元素的并行重分配处理

这会导致大部分元素未被正确重分配到buffer中，直接引发排序结果错误。

3. 潜在的线程局部偏移量竞争问题 #

虽然localOffset是每个线程独立的，但在重分配阶段如果未正确按元素bucket对应到localOffset的位置自增，也可能导致元素覆盖或位置错误。

修复建议 #

1. 修复剩余元素的直方图统计逻辑 #

将剩余元素的处理纳入OpenMP parallel区域，让每个线程仅处理自己负责的chunk内的剩余元素，确保所有线程的直方图都能正确统计对应元素：

// Step1: 直方图计数（修复后）
#pragma omp parallel
{
    auto &hist = localHist[omp_get_thread_num()];
    size_t tid = omp_get_thread_num();
    size_t nThreads = omp_get_num_threads();
    size_t chunkSize = (n + nThreads - 1) / nThreads;
    size_t start = tid * chunkSize;
    size_t end = std::min(start + chunkSize, n);

    // 处理当前chunk内的向量元素（i+=4）
    for (size_t i = start; i <= end - 4; i += 4) {
        __m256i vec = _mm256_loadu_si256((__m256i*)&keys[i]);
        __m256i shifted_right = _mm256_srli_epi64(vec, shift);
        __m256i result_and = _mm256_and_si256(shifted_right, mask_vec);
        alignas(32) uint64_t result[4];
        _mm256_storeu_si256((__m256i *)result, result_and);
        hist[result[0]]++;
        hist[result[1]]++;
        hist[result[2]]++;
        hist[result[3]]++;
    }

    // 处理当前chunk内的剩余元素
    for (size_t i = (end / 4) * 4; i < end; ++i) {
        size_t bucket = (keys[i] >> shift) & BUCKET_MASK;
        hist[bucket]++;
    }
}

2. 补全重分配阶段的AVX代码及剩余元素处理 #

确保重分配阶段正确计算每个元素的存储位置，并将元素写入buffer的对应位置，同时并行处理剩余元素：

// Step3: 元素重分配（修复后）
#pragma omp parallel
{
    auto &localOffset = localHist[omp_get_thread_num()];
    size_t tid = omp_get_thread_num();
    size_t nThreads = omp_get_num_threads();
    size_t chunkSize = (n + nThreads - 1) / nThreads;
    size_t start = tid * chunkSize;
    size_t end = std::min(start + chunkSize, n);

    // 处理当前chunk内的向量元素
    for (size_t i = start; i <= end - 4; i += 4) {
        __m256i vec = _mm256_loadu_si256((__m256i*)&keys[i]);
        __m256i shifted_right = _mm256_srli_epi64(vec, shift);
        __m256i result_and = _mm256_and_si256(shifted_right, mask_vec);
        alignas(32) uint64_t buckets[4];
        _mm256_storeu_si256((__m256i *)buckets, result_and);

        // 获取每个元素的存储位置并更新偏移
        size_t pos0 = localOffset[buckets[0]]++;
        size_t pos1 = localOffset[buckets[1]]++;
        size_t pos2 = localOffset[buckets[2]]++;
        size_t pos3 = localOffset[buckets[3]]++;

        // 写入keys
        buffer[pos0] = keys[i];
        buffer[pos1] = keys[i+1];
        buffer[pos2] = keys[i+2];
        buffer[pos3] = keys[i+3];
        // 写入解码后的点
        bufferDecoded[pos0] = points[i];
        bufferDecoded[pos1] = points[i+1];
        bufferDecoded[pos2] = points[i+2];
        bufferDecoded[pos3] = points[i+3];
        // 写入元数据（如果存在）
        if (meta_opt) {
            metadata_buffer[pos0] = (*meta_opt)[i];
            metadata_buffer[pos1] = (*meta_opt)[i+1];
            metadata_buffer[pos2] = (*meta_opt)[i+2];
            metadata_buffer[pos3] = (*meta_opt)[i+3];
        }
    }

    // 处理当前chunk内的剩余元素
    for (size_t i = (end /4)*4; i < end; ++i) {
        size_t bucket = (keys[i] >> shift) & BUCKET_MASK;
        size_t pos = localOffset[bucket]++;
        buffer[pos] = keys[i];
        bufferDecoded[pos] = points[i];
        if (meta_opt) {
            metadata_buffer[pos] = (*meta_opt)[i];
        }
    }
}

内容来源于stack exchange