沐曦的__ballot_sync应该是有问题

/*
 * 测试用例：演示沐曦GPU __ballot_sync bug及workaround验证
 * 
 * 问题：在warp 1中，线程59-63(lane_id 27-31)执行__ballot_sync时，
 *       即使predicate为真，返回0x00000000而不是正确的bitmask
 * 
 */

#include <mc_runtime.h>
#include <stdio.h>

#define USE_METAX  // 启用workaround


// 原始的__ballot_sync实现（NVIDIA GPU正常，沐曦GPU有bug）
__device__ uint32_t ballot_original(bool pred) {
    uint32_t FULL_WARP_ACTIVE_BMSK = 0xFFFFFFFF;
    return __ballot_sync(FULL_WARP_ACTIVE_BMSK, pred);
}

// Workaround：使用共享内存实现ballot（沐曦GPU专用）
__device__ uint32_t ballot_workaround(bool pred) {
    uint32_t lane_id = threadIdx.x % 32;
    uint32_t thrdIdxInBlk = threadIdx.x + threadIdx.y * blockDim.x;
    uint32_t warp_id = thrdIdxInBlk / 32;

    __shared__ uint32_t s_warp_pred[512];

    uint32_t idx = warp_id * 32 + lane_id;
    s_warp_pred[idx] = pred ? 1 : 0;
    __syncthreads();

    uint32_t ballotResult = 0;
    for (int i = 0; i < 32; i++) {
        if (s_warp_pred[warp_id * 32 + i]) {
            ballotResult |= (1U << i);
        }
    }
    __syncthreads();

    return ballotResult;
}

// 测试内核：重现生产环境中的bug场景
__global__ void test_ballot_bug(
    uint32_t* d_original_results,
    uint32_t* d_workaround_results,
    uint32_t nInfoBits
) {
    uint32_t thrdIdxInBlk = threadIdx.x + threadIdx.y * blockDim.x;
    uint32_t lane_id = threadIdx.x % 32;
    uint32_t warp_id = thrdIdxInBlk / 32;

    // 生产环境中的精确predicate逻辑
    uint32_t N_MAX_INFO_BITS = 128;
    int16_t interleaverTblStartIdx = N_MAX_INFO_BITS - nInfoBits;
    bool pred = ((thrdIdxInBlk < N_MAX_INFO_BITS) && (thrdIdxInBlk >= interleaverTblStartIdx)) ? true : false;

    // 测试原始实现
    uint32_t original_ballot = ballot_original(pred);

    // 测试workaround
    uint32_t workaround_ballot = ballot_workaround(pred);

    // 保存结果（每个warp只保存第一个线程的结果）
    if (lane_id == 0 && warp_id < 2) {  // 只保存warp 0和warp 1
        d_original_results[warp_id] = original_ballot;
        d_workaround_results[warp_id] = workaround_ballot;
    }

    // 打印调试信息
    if (pred) {
        printf("threadIdx.x=%d, threadIdx.y=%d, thrdIdxInBlk=%d, lane_id=%d, warp_id=%d, pred=%d\n",
               threadIdx.x, threadIdx.y, thrdIdxInBlk, lane_id, warp_id, pred);
        printf("  original_ballot=0x%08x, workaround_ballot=0x%08x\n",
               original_ballot, workaround_ballot);
    }
}

// 验证结果的辅助函数
void verify_results(uint32_t* original, uint32_t* workaround, int nWarps, uint32_t expected_mask) {
    printf("\n========== 验证结果 ==========\n");
    printf("预期bitmask (所有pred为true的线程): 0x%08x\n", expected_mask);
    printf("\n原始 __ballot_sync 结果:\n");
    for (int i = 0; i < nWarps; i++) {
        printf("  Warp %d: 0x%08x %s\n", i, original[i],
               (original[i] == expected_mask) ? "✓ 正确" : "✗ 错误");
    }

    printf("\nWorkaround 结果:\n");
    for (int i = 0; i < nWarps; i++) {
        printf("  Warp %d: 0x%08x %s\n", i, workaround[i],
               (workaround[i] == expected_mask) ? "✓ 正确" : "✗ 错误");
    }
}

int main() {
    // 测试参数：重现生产环境中的bug场景
    uint32_t nInfoBits = 69;  // 导致线程59-63的predicate为true
    uint32_t N_MAX_INFO_BITS = 128;

    printf("========== 测试配置 ==========\n");
    printf("nInfoBits = %d, N_MAX_INFO_BITS = %d\n", nInfoBits, N_MAX_INFO_BITS);
    printf("Block dim: 32x16x1 (512 threads)\n");
    printf("Grid dim: 1x1x1\n\n");

    // 计算预期的bitmask
    // 线程59-63 (lane 27-31 in warp 1) 应该有 pred=true
    // 预期bitmask = 0xF8000000 (lanes 27-31 set)
    uint32_t expected_mask = 0xF8000000;

    // 分配设备内存
    uint32_t *d_original, *d_workaround;
    uint32_t *h_original, *h_workaround;

    mcMalloc(&d_original, 2 * sizeof(uint32_t));
    mcMalloc(&d_workaround, 2 * sizeof(uint32_t));
    h_original = (uint32_t*)malloc(2 * sizeof(uint32_t));
    h_workaround = (uint32_t*)malloc(2 * sizeof(uint32_t));

    // 初始化
    mcMemset(d_original, 0, 2 * sizeof(uint32_t));
    mcMemset(d_workaround, 0, 2 * sizeof(uint32_t));

    // 启动kernel
    printf("========== 启动Kernel ==========\n");
    dim3 blockDim(32, 16, 1);  // 32x16x1 = 512 threads
    dim3 gridDim(1, 1, 1);

    test_ballot_bug<<<gridDim, blockDim>>>(d_original, d_workaround, nInfoBits);

    // 同步并检查错误
    mcError_t err = mcDeviceSynchronize();
    if (err != mcSuccess) {
        printf("CUDA Error: %s\n", mcGetErrorString(err));
        return -1;
    }

    // 复制结果回主机
    mcMemcpy(h_original, d_original, 2 * sizeof(uint32_t), mcMemcpyDeviceToHost);
    mcMemcpy(h_workaround, d_workaround, 2 * sizeof(uint32_t), mcMemcpyDeviceToHost);

    // 验证结果
    verify_results(h_original, h_workaround, 2, expected_mask);

    // 清理
    mcFree(d_original);
    mcFree(d_workaround);
    free(h_original);
    free(h_workaround);

    printf("\n========== 测试完成 ==========\n");
    return 0;
}
Warp 0, 2, 3的__ballot_sync工作正常
只有Warp 1的特定lane范围(27-31)受影响
Warp 1的lanes 27-31 (线程59-63)在predicate=true时，__ballot_sync返回0x00000000