CUDA编程：Histogram 直方图统计与 Atomic 优化

1. 什么是 histogram 直方图
2. 为什么 histogram 容易产生 atomic 冲突
3. 掌握 global atomic histogram 的基础写法
4. 掌握 shared memory 局部 histogram 优化
5. 理解“先局部统计，再全局合并”的优化思想
6. 对比不同 bin 数量、数据分布下的性能差异

input = [1, 2, 1, 3, 2, 1, 0, 3]

hist[0] = 1
hist[1] = 3
hist[2] = 2
hist[3] = 2

hist = [1, 3, 2, 2]

int bin = input[idx];
atomicAdd(&hist[bin], 1);

input = [1, 1, 1, 1, 1, 1, ...]

atomicAdd(&hist[1], 1);

多个线程可能同时写同一个地址
必须用 atomic 保证正确性
但 atomic 冲突会影响性能

atomicAdd(&global_hist[bin], 1);

代码简单
结果正确
适合 bin 多、冲突不严重的情况

所有 block 都直接竞争 global memory
当数据集中在少数 bin 时会很慢

每个 block 先在 shared memory 中统计一个局部 hist
然后每个 block 再把局部 hist 合并到 global hist

Global input
    ↓
每个 block 建立 shared_hist
    ↓
block 内线程 atomicAdd(shared_hist[bin], 1)
    ↓
block 内统计完成
    ↓
把 shared_hist 合并到 global_hist

block 内 shared atomic
+
少量 global atomic

1. CPU histogram
2. GPU global atomic histogram
3. GPU shared memory histogram

1. bin 数量：16、64、256
2. 数据规模：1M、16M、64M
3. 数据分布：
   - 均匀分布：冲突较平均
   - 热点分布：大量数据集中在少数 bin

lesson12_histogram.cu

#include <cuda_runtime.h>

#include <algorithm>
#include <chrono>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <random>
#include <vector>

#define CUDA_CHECK(call)                                                        \
    do {                                                                        \
        cudaError_t err = call;                                                 \
        if (err != cudaSuccess) {                                               \
            std::cerr << "CUDA Error: " << cudaGetErrorString(err)              \
                      << " at " << __FILE__ << ":" << __LINE__ << std::endl;    \
            std::exit(EXIT_FAILURE);                                            \
        }                                                                       \
    } while (0)

/*
 * CPU histogram reference.
 */
void histogram_cpu(const int* input, int* hist, size_t n, int bins) {
    std::fill(hist, hist + bins, 0);

    for (size_t i = 0; i < n; ++i) {
        int bin = input[i];
        hist[bin]++;
    }
}

/*
 * 方法 1：Global atomic histogram。
 *
 * 每个线程读取一个 input，然后直接 atomicAdd 到 global hist。
 */
__global__ void histogram_global_atomic_kernel(const int* input,
                                               int* global_hist,
                                               size_t n) {
    size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;

    if (idx < n) {
        int bin = input[idx];
        atomicAdd(&global_hist[bin], 1);
    }
}

/*
 * 方法 2：Shared memory 局部 histogram。
 *
 * 每个 block 先在 shared memory 中统计局部 hist，
 * 然后再合并到 global hist。
 */
__global__ void histogram_shared_kernel(const int* input,
                                        int* global_hist,
                                        size_t n,
                                        int bins) {
    extern __shared__ int shared_hist[];

    unsigned int tid = threadIdx.x;

    /*
     * 1. 初始化 shared_hist。
     *
     * bins 可能大于 blockDim.x，所以用 stride 循环。
     */
    for (int i = tid; i < bins; i += blockDim.x) {
        shared_hist[i] = 0;
    }

    __syncthreads();

    /*
     * 2. 每个线程处理多个 input 元素。
     *
     * grid-stride loop：
     * 让总线程数不足以覆盖 n 时，也可以处理完整数组。
     */
    size_t global_tid = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    size_t stride = static_cast<size_t>(blockDim.x) * gridDim.x;

    for (size_t i = global_tid; i < n; i += stride) {
        int bin = input[i];
        atomicAdd(&shared_hist[bin], 1);
    }

    __syncthreads();

    /*
     * 3. 把 shared_hist 合并到 global_hist。
     *
     * 每个 block 对每个 bin 最多做一次 global atomic。
     */
    for (int i = tid; i < bins; i += blockDim.x) {
        if (shared_hist[i] > 0) {
            atomicAdd(&global_hist[i], shared_hist[i]);
        }
    }
}

/*
 * 用 cudaEvent 计时一个 kernel。
 */
template <typename LaunchFunc>
float time_kernel(LaunchFunc launch, int repeat) {
    /*
     * warmup
     */
    launch();
    CUDA_CHECK(cudaDeviceSynchronize());

    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));

    float total_ms = 0.0f;

    for (int r = 0; r < repeat; ++r) {
        CUDA_CHECK(cudaEventRecord(start));

        launch();

        CUDA_CHECK(cudaEventRecord(stop));
        CUDA_CHECK(cudaEventSynchronize(stop));

        float ms = 0.0f;
        CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
        total_ms += ms;
    }

    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));

    return total_ms / repeat;
}

/*
 * 校验两个 histogram 是否一致。
 */
bool check_hist(const std::vector<int>& a,
                const std::vector<int>& b,
                int bins) {
    for (int i = 0; i < bins; ++i) {
        if (a[i] != b[i]) {
            std::cerr << "Mismatch at bin " << i
                      << ": CPU=" << a[i]
                      << ", GPU=" << b[i] << std::endl;
            return false;
        }
    }

    return true;
}

/*
 * 生成输入数据。
 *
 * mode = uniform:
 *   每个 bin 概率基本相同。
 *
 * mode = hot:
 *   90% 的数据集中在 bin 0，剩下 10% 随机分布。
 */
void generate_input(std::vector<int>& input,
                    int bins,
                    const std::string& mode) {
    std::mt19937 rng(12345);
    std::uniform_int_distribution<int> uniform_dist(0, bins - 1);
    std::uniform_real_distribution<float> prob(0.0f, 1.0f);

    for (size_t i = 0; i < input.size(); ++i) {
        if (mode == "hot") {
            if (prob(rng) < 0.90f) {
                input[i] = 0;
            } else {
                input[i] = uniform_dist(rng);
            }
        } else {
            input[i] = uniform_dist(rng);
        }
    }
}

int main(int argc, char** argv) {
    size_t n = 1ULL << 24;     // 16,777,216 elements
    int bins = 256;
    int block_size = 256;
    int repeat = 10;
    std::string mode = "uniform";

    if (argc >= 2) {
        n = static_cast<size_t>(std::atoll(argv[1]));
    }

    if (argc >= 3) {
        bins = std::atoi(argv[2]);
    }

    if (argc >= 4) {
        block_size = std::atoi(argv[3]);
    }

    if (argc >= 5) {
        repeat = std::atoi(argv[4]);
    }

    if (argc >= 6) {
        mode = argv[5];
    }

    if (bins <= 0 || bins > 4096) {
        std::cerr << "bins must be in (0, 4096].\n";
        return 1;
    }

    if (block_size <= 0 || block_size > 1024) {
        std::cerr << "block_size must be in (0, 1024].\n";
        return 1;
    }

    size_t input_bytes = n * sizeof(int);
    size_t hist_bytes = static_cast<size_t>(bins) * sizeof(int);

    std::cout << "CUDA Lesson 12: Histogram and Atomic Optimization\n";
    std::cout << "Elements   : " << n << "\n";
    std::cout << "Input size : " << input_bytes / 1024.0 / 1024.0 << " MB\n";
    std::cout << "Bins       : " << bins << "\n";
    std::cout << "Block size : " << block_size << "\n";
    std::cout << "Repeat     : " << repeat << "\n";
    std::cout << "Mode       : " << mode << "\n";

    std::vector<int> h_input(n);
    std::vector<int> h_cpu_hist(bins, 0);
    std::vector<int> h_global_hist(bins, 0);
    std::vector<int> h_shared_hist(bins, 0);

    generate_input(h_input, bins, mode);

    /*
     * CPU timing.
     */
    auto cpu_start = std::chrono::high_resolution_clock::now();

    histogram_cpu(h_input.data(),
                  h_cpu_hist.data(),
                  n,
                  bins);

    auto cpu_end = std::chrono::high_resolution_clock::now();

    double cpu_ms = std::chrono::duration<double, std::milli>(
        cpu_end - cpu_start
    ).count();

    int* d_input = nullptr;
    int* d_global_hist = nullptr;
    int* d_shared_hist = nullptr;

    CUDA_CHECK(cudaMalloc(&d_input, input_bytes));
    CUDA_CHECK(cudaMalloc(&d_global_hist, hist_bytes));
    CUDA_CHECK(cudaMalloc(&d_shared_hist, hist_bytes));

    CUDA_CHECK(cudaMemcpy(d_input,
                          h_input.data(),
                          input_bytes,
                          cudaMemcpyHostToDevice));

    /*
     * Global atomic kernel 配置。
     */
    int grid_global = static_cast<int>((n + block_size - 1) / block_size);

    auto launch_global = [&]() {
        CUDA_CHECK(cudaMemset(d_global_hist, 0, hist_bytes));

        histogram_global_atomic_kernel<<<grid_global, block_size>>>(
            d_input,
            d_global_hist,
            n
        );
        CUDA_CHECK(cudaGetLastError());
    };

    /*
     * Shared histogram kernel 配置。
     *
     * 使用固定 block 数量，配合 grid-stride loop。
     * 不要让 block 数过大，否则每个 block 都要合并一次 hist，global atomic 次数也会上升。
     */
    int sm_count = 0;
    CUDA_CHECK(cudaDeviceGetAttribute(&sm_count,
                                      cudaDevAttrMultiProcessorCount,
                                      0));

    int grid_shared = sm_count * 4;
    size_t shared_bytes = hist_bytes;

    auto launch_shared = [&]() {
        CUDA_CHECK(cudaMemset(d_shared_hist, 0, hist_bytes));

        histogram_shared_kernel<<<grid_shared,
                                  block_size,
                                  shared_bytes>>>(
            d_input,
            d_shared_hist,
            n,
            bins
        );
        CUDA_CHECK(cudaGetLastError());
    };

    float global_ms = time_kernel(launch_global, repeat);
    float shared_ms = time_kernel(launch_shared, repeat);

    CUDA_CHECK(cudaMemcpy(h_global_hist.data(),
                          d_global_hist,
                          hist_bytes,
                          cudaMemcpyDeviceToHost));

    CUDA_CHECK(cudaMemcpy(h_shared_hist.data(),
                          d_shared_hist,
                          hist_bytes,
                          cudaMemcpyDeviceToHost));

    bool global_ok = check_hist(h_cpu_hist, h_global_hist, bins);
    bool shared_ok = check_hist(h_cpu_hist, h_shared_hist, bins);

    std::cout << std::fixed << std::setprecision(4);

    std::cout << "\n[Timing]\n";
    std::cout << "CPU histogram time          : " << cpu_ms << " ms\n";
    std::cout << "GPU global atomic time      : " << global_ms << " ms\n";
    std::cout << "GPU shared histogram time   : " << shared_ms << " ms\n";

    std::cout << "\n[Speedup]\n";
    std::cout << "Global atomic vs CPU        : " << cpu_ms / global_ms << "x\n";
    std::cout << "Shared histogram vs CPU     : " << cpu_ms / shared_ms << "x\n";
    std::cout << "Shared vs Global atomic     : " << global_ms / shared_ms << "x\n";

    std::cout << "\n[Check]\n";
    std::cout << "Global atomic result        : " << (global_ok ? "PASS" : "FAIL") << "\n";
    std::cout << "Shared histogram result     : " << (shared_ok ? "PASS" : "FAIL") << "\n";

    std::cout << "\n[Sample hist bins]\n";
    int show_bins = std::min(bins, 16);
    for (int i = 0; i < show_bins; ++i) {
        std::cout << "bin[" << i << "] = " << h_cpu_hist[i] << "\n";
    }

    CUDA_CHECK(cudaFree(d_input));
    CUDA_CHECK(cudaFree(d_global_hist));
    CUDA_CHECK(cudaFree(d_shared_hist));

    return 0;
}

nvcc -O3 -arch=sm_75 lesson12_histogram.cu -o lesson12_histogram

./lesson12_histogram

./lesson12_histogram 16777216 256 256 10 uniform

第 1 个参数：元素数量 n
第 2 个参数：bins 数量
第 3 个参数：block_size
第 4 个参数：repeat
第 5 个参数：mode，uniform 或 hot

./lesson12_histogram 16777216 16 256 10 uniform
./lesson12_histogram 16777216 64 256 10 uniform
./lesson12_histogram 16777216 256 256 10 uniform

./lesson12_histogram 16777216 16 256 10 hot
./lesson12_histogram 16777216 64 256 10 hot
./lesson12_histogram 16777216 256 256 10 hot

90% 数据集中在 bin 0
global atomic 明显变慢
shared histogram 通常更有优势

./lesson12_histogram 1048576 256 256 10 uniform
./lesson12_histogram 16777216 256 256 10 uniform
./lesson12_histogram 67108864 256 256 10 uniform

数据越大，GPU 并行统计优势越明显

1. Histogram 是典型的 many-to-few 统计问题
2. 多线程更新同一个 bin 时必须处理写冲突
3. global atomic 写法简单，但热点 bin 会造成严重竞争
4. shared memory 局部 histogram 可以减少 global atomic 次数
5. 优化思想是“先 block 内局部统计，再跨 block 合并”
6. bins 数量和数据分布会显著影响性能