寻找256位整数中第k个置位的索引及代码优化需求

问题说明 #

需要实现一个函数，从由4个uint64_t组成的256位整数中，找到第k个置位（1）的位索引（范围0-255，最低位为0）。原代码逻辑存在偏差，且第二个循环效率较低，希望通过位操作指令（如clz、bsf）以及SIMD指令（SSE/AVX）优化性能，可使用GCC/Clang/Linux专属扩展，不考虑可移植性。

原代码分析 #

原代码如下：

uint16_t find_index(const uint64_t A[4], uint8_t k) {
    uint16_t res = 0;
    for (int i = 0; i < (k / 64) - 1; i++) {
        res += __builtin_popcountll(A[i]);
    }
  
    for (int i = 0; i < k % 64; i++) {
        res += (A[k / 64] >> ((uint64_t)i)) & 0x01;
    }

    return res;
}

原代码存在两个核心问题：

1. 逻辑偏差：当前代码计算的是前(k/64)*64 + (k%64)位中的置位总数，而非第k个置位的位索引。
2. 效率低下：第二个循环逐位检查目标块的前k%64位，时间复杂度为O(64)，可通过位操作指令优化为O(1)。

优化方案1：修正逻辑并使用位操作指令优化 #

首先修正核心逻辑：先逐个统计每个64位块的置位数量，定位到包含第k个置位的块；再在该块内找到目标置位的位置，最后计算其在256位中的总索引。

利用GCC内置函数和BMI2指令（需CPU支持，如Intel Haswell、AMD Ryzen及以后）优化块内查找：

#include <stdint.h>
#include <x86intrin.h>

uint16_t find_index_bmi2(const uint64_t A[4], uint8_t k) {
    uint16_t total = 0;
    int block_idx;

    // 定位包含第k个置位的块
    for (block_idx = 0; block_idx < 4; block_idx++) {
        uint16_t cnt = __builtin_popcountll(A[block_idx]);
        if (total + cnt >= k) {
            break;
        }
        total += cnt;
    }

    // 在目标块内快速定位第(k - total)个置位
    uint64_t block = A[block_idx];
    uint8_t target_in_block = k - total;
    // 构造仅第target_in_block位为1的掩码
    uint64_t mask = 1ULL << (target_in_block - 1);
    // 用_pdep_u64将掩码映射到块的置位位置，得到唯一的目标置位
    uint64_t res_block = _pdep_u64(mask, block);
    // 获取该置位的位偏移
    int bit_pos = __builtin_ctzll(res_block);

    // 计算256位中的总索引
    return block_idx * 64 + bit_pos;
}

此方案中，_pdep_u64指令将块内查找从O(n)降为O(1)，__builtin_ctzll快速获取置位的偏移位置，性能大幅提升。

优化方案2：使用AVX2指令加速块内置位统计 #

对于256位整数，可利用AVX2指令一次性统计所有4个64位块的置位数量，减少循环开销：

#include <stdint.h>
#include <x86intrin.h>

uint16_t find_index_avx2_bmi2(const uint64_t A[4], uint8_t k) {
    // 将4个uint64_t加载到AVX2寄存器
    __m256i vec = _mm256_loadu_si256((const __m256i*)A);
    // 一次性统计4个64位元素的置位数量
    __m256i popcounts = _mm256_popcnt_epi64(vec);
    
    // 将统计结果存储到数组
    uint32_t counts[4];
    _mm256_storeu_si256((__m256i*)counts, popcounts);

    uint16_t total = 0;
    int block_idx;
    // 定位目标块
    for (block_idx = 0; block_idx < 4; block_idx++) {
        if (total + counts[block_idx] >= k) {
            break;
        }
        total += counts[block_idx];
    }

    // 复用BMI2的块内查找逻辑
    uint64_t block = A[block_idx];
    uint8_t target_in_block = k - total;
    uint64_t mask = 1ULL << (target_in_block - 1);
    uint64_t res_block = _pdep_u64(mask, block);
    int bit_pos = __builtin_ctzll(res_block);

    return block_idx * 64 + bit_pos;
}

AVX2的_mm256_popcnt_epi64指令批量完成置位统计，相比逐个调用__builtin_popcountll进一步降低了循环开销。

编译注意事项 #

• 使用BMI2指令需添加编译选项-mbmi2
• 使用AVX2指令需添加编译选项-mavx2
• 配合-O3选项可让编译器进一步优化代码逻辑

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