cub::BlockHistogram 如何处理超额线程对应的无效数据？

2026-11-04数据小许

可以增加一个 bin 来接收无效样本，统计后只保留原来的 256 个 bin。如果必须维持 256 个 bin，也可以将无效位置填零，再减去补齐数量。对于连续数据尾部的补齐，数量可以直接计算，不必逐个统计无效样本。

已核对的官方 Histogram 和 Composite 接口没有 valid_items 参数，并且要求所有输入满足 0 <= input[i] < Bins。因此，不能直接给原来的 256-bin 配置传入 256 或其他越界值来表示“忽略”。具体接口仍应以项目使用的 CUB 版本为准。官方 API 文档

用额外的 bin 接收无效样本

假设有效样本都是 [0, 256) 范围内的整数，可以把 Bins 改成 257，并将无效位置填为 256。无效样本会进入 histogram[256]，输出时只需复制前 256 个 bin。

在 257-bin 配置下，256 是合法的 bin 索引，原来的 std::uint16_t 类型也能表示这个值。

下面的 kernel 使用一维 block，每个 block 固定为 256 个线程，每个线程处理 ItemsPerThread 个样本：

#include <cub/block/block_histogram.cuh>
#include <cstddef>
#include <cstdint>

template <int ItemsPerThread>
__global__ void histogram_per_block(
    const std::uint16_t* input,
    std::size_t n,
    unsigned int* block_histograms)
{
    static_assert(ItemsPerThread > 0, "ItemsPerThread must be positive");

    constexpr int Threads = 256;
    constexpr int OutputBins = 256;
    constexpr int TotalBins = OutputBins + 1;
    constexpr std::uint16_t InvalidBin = 256;
    constexpr std::size_t TileItems =
        std::size_t{Threads} * ItemsPerThread;

    using BlockHistogram = cub::BlockHistogram<
        std::uint16_t,
        Threads,
        ItemsPerThread,
        TotalBins,
        cub::BLOCK_HISTO_ATOMIC>;

    __shared__ typename BlockHistogram::TempStorage temp_storage;
    __shared__ unsigned int histogram[TotalBins];

    std::uint16_t items[ItemsPerThread];
    const std::size_t base = std::size_t{blockIdx.x} * TileItems;

    #pragma unroll
    for (int i = 0; i < ItemsPerThread; ++i) {
        const std::size_t index =
            base + threadIdx.x + std::size_t{i} * Threads;
        items[i] = (index < n) ? input[index] : InvalidBin;
    }

    BlockHistogram(temp_storage).Histogram(items, histogram);
    __syncthreads();

    block_histograms[std::size_t{blockIdx.x} * OutputBins
                     + threadIdx.x] = histogram[threadIdx.x];
}

使用时需要满足以下条件：

以 <<<num_blocks, 256>>> 启动，block 的形状与模板配置一致。
num_blocks 至少覆盖全部数据；每个 block 最多处理 256 * ItemsPerThread 个样本。
block_histograms 至少容纳 num_blocks * 256 个计数。输出的是各 block 的局部直方图，整体结果还需要按 bin 合并。
有效输入必须位于 [0, 256)。std::uint16_t 是存储类型，其全部取值不能直接用于这个配置。
n == 0 时可以跳过启动，并按调用方需要将最终结果置零。

代码按每个样本的索引判断有效性，因此也能处理一个线程只有部分样本有效的情况。

保留 256 个 bin 时，直接计算补齐数量

如果希望保留原来的模板配置，可以继续把无效位置填为 0，随后修正第 0 个 bin。

每个 block 处理一段连续数据时，设该段起始索引为 base：

capacity = 256 * items_per_thread
valid = base < N ? min(capacity, N - base) : 0
padding = capacity - valid

所有线程完成 Histogram 后，再由一个线程减去 padding：

BlockHistogram(temp_storage).Histogram(items, histogram);
__syncthreads();

if (threadIdx.x == 0) {
    histogram[0] -= static_cast<unsigned int>(padding);
}

__syncthreads();

第二个同步点保证其他线程读取或输出直方图时，修正已经完成。减去的只是人为补齐的数量，真实的零值样本仍会保留。计数类型需要能够容纳补齐前的计数。

如果无效位置来自任意分布的有效性掩码，而不只是数组尾部，就不能直接套用上述公式。此时需要统计掩码中的无效数量，或使用额外 bin 的方式。

为什么不能只让有效线程调用 Histogram？

BlockHistogram 是 block 范围的集体操作。不要用下面的方式跳过无效线程：

// 不应这样调用 block 范围的集体操作。
if (index < n) {
    BlockHistogram(temp_storage).Histogram(items, histogram);
}

官方实现中的 Histogram 会先初始化直方图，再执行 __syncthreads()，然后统计每个线程提供的样本。应让 block 中所有线程一致地参与调用，只在加载样本时判断有效性。__activemask() 返回的是 warp 的活动线程掩码，不能将这个接口变成支持部分线程参与的 block 操作。BlockHistogram 实现

额外 bin 仍会执行无效样本对应的更新，不能因此断言它比原有的条件统计更快。尾部补齐占比很高，或大量样本集中在同一个 bin 时，应结合实际数据比较性能。