LeetCUDA学习记录（二）——histogram算子（直方图统计）代码实现详细解读
🎓 一、环境配置安装
🌵 1.1 环境配置和安装
有关环境配置和安装，以及具体gpu（如3090）上的编译错误问题，请参考博客：LeetCUDA学习记录（一）——elementwise算子代码实现详细解读
🍀 1.2 histogram算子运行体验
简介：histogram算子通过PyTorch C++/CUDA 扩展实现 GPU 加速的直方图统计。对应仓库难度：Easy⭐️
直方图统计是一种将数据按区间（或离散值）进行分组，并统计每个组内数据出现次数的方法。它可以直观地展示数据的分布特征。例如，对于包含重复整数的序列[0,1,2,…,9,0,1,2,…,9,…]，直方图统计会计算每个整数（0-9）出现的总次数。在图像处理、数据挖掘等领域，直方图统计是基础且高频的操作，GPU 加速能显著提升其处理大规模数据的效率。
笔者使用的是3090显卡，因此需修改histogram.py中的配置，在extra_cuda_cflags中添加以下项：
# new add, 添加针对不同 GPU 架构的编译选项
"-gencode=arch=compute_86,code=sm_86", # 3090
按照上面配置重新运行histogram.py程序：
CUDA_VISIBLE_DEVICES=1 python3 histogram.py
得到结果如下（更多结果已省略），可以看出已经正确执行：
🚀 二、histogram算子代码解读
🏅 2.1 histogram算子代码结构
整个LeetCUDA histogram算子包括三部分文件：
histogram.py: 用 PyTorch 调用自定义的 CUDA 核函数，在 GPU 上实现两种不同方式的直方图统计 —— 一种是基础的单元素处理，另一种是用 SIMD 技术一次处理 4 个元素的优化版本。通过生成包含 0-9 重复 1000
次的测试数据，分别调用两种方法计算每个数字的出现次数。histogram.cu: 实现两种不同方式的直方图统计的 CUDA 内核，并通过 PyTorch C++ 扩展暴露给 Python 使用，用于实现高性能直方图统计功能。README.md: 代码的运行说明和运行结果。
🀄️ 2.2 histogram.py代码解读
🏆 2.2.1 导入模块与设置
torch.utils.cpp_extension.load：动态编译并加载 CUDA / C++ 扩展。torch.set_grad_enabled(False)：关闭梯度计算，减少开销，提高性能。
import torch
from torch.utils.cpp_extension import load
torch.set_grad_enabled(False) # 全局关闭梯度计算（只做推理/性能测试时更快）
🌟 2.2.2 加载 CUDA 内核扩展
使用 torch.utils.cpp_extension.load 将 CUDA 内核编译成 Python 可调用模块。支持不同版本直方图统计计算的内核。
-gencode指定 GPU 架构，以确保内核可以在对应设备上运行。--use_fast_math启用快速数学运算，牺牲精度换性能。
# Load the CUDA kernel as a python module
lib = load(
name="hist_lib",
sources=["histogram.cu"], # CUDA 源文件
extra_cuda_cflags=[
"-O3",
"-U__CUDA_NO_HALF_OPERATORS__", # 启用半精度 float16 运算支持
"-U__CUDA_NO_HALF_CONVERSIONS__",
"-U__CUDA_NO_HALF2_OPERATORS__",
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
"--expt-relaxed-constexpr",
"--expt-extended-lambda", # 允许 CUDA lambda 扩展
"--use_fast_math", # 使用快速数学运算（牺牲精度换速度）
# new add, 添加针对不同 GPU 架构的编译选项
"-gencode=arch=compute_86,code=sm_86" # 3090
],
extra_cflags=["-std=c++17"],
)
🍎 2.2.3 测试不同方法实现的算子的正确性
这段代码主要是通过生成包含 0-9 重复 1000 次的测试数据，分别调用两种方法计算每个数字的出现次数，最后打印结果来验证优化版在保证统计正确的同时，是否能利用 GPU 的并行计算能力提升效率。简单说，就是用代码验证 “GPU 加速 + 向量计算优化” 能不能让直方图统计又快又准。
# 1、构造测试数据：生成包含1000次重复的0-9整数序列，类型为int32，并转移到GPU
# 预期每个整数（0-9）的出现次数均为1000
a = torch.tensor(list(range(10)) * 1000, dtype=torch.int32).cuda()
# 2、调用基础版CUDA核函数计算直方图（单元素处理）
h_i32 = lib.histogram_i32(a)
print("-" * 80)
# 打印基础版结果，验证每个整数的统计次数
for i in range(h_i32.shape[0]):
print(f"h_i32 {i}: {h_i32[i]}")
print("-" * 80)
# 3、调用SIMD优化版CUDA核函数计算直方图（4元素并行处理，利用向量计算）
h_i32x4 = lib.histogram_i32x4(a)
# 打印优化版结果，对比验证统计正确性
for i in range(h_i32x4.shape[0]):
print(f"h_i32x4 {i}: {h_i32x4[i]}")
print("-" * 80)
🌕 2.3 histogram.cu代码解读
⚡️ 2.3.1 头文件和宏定义
这段代码引入 CUDA 运行时库和 PyTorch 扩展接口。同时定义了针对向量计算的优化宏 ——INT4和FLOAT4，能将连续内存中的 4 个整数或浮点数转换为 CUDA 原生向量类型，配合 GPU 的 SIMD（单指令多数据）计算能力实现并行加速。同时，WARP_SIZE的定义适配了 NVIDIA GPU 的硬件特性（32 线程为一个基本执行单元），为后续核函数中线程分工、内存访问优化打下基础。
#include <algorithm> // 引入标准库头文件，提供算法、内存管理等基础功能
#include <cuda_runtime.h> // CUDA运行时API头文件，用于CUDA设备管理、内存分配等底层操作
#include <float.h> // 定义浮点数相关常量（如FLT_MAX）的头文件
#include <stdio.h> // 标准输入输出头文件，用于调试打印等
#include <stdlib.h> // 标准库头文件，提供动态内存分配（malloc等）功能
#include <torch/extension.h> // PyTorch C++扩展头文件，用于与PyTorch张量交互、定义扩展函数
#include <torch/types.h> // PyTorch类型定义头文件，包含张量类型等定义
#include <tuple> // 提供tuple容器支持，用于返回多值结果
#include <vector> // 提供向量容器支持，用于动态数组操作
// 定义GPU warp大小为32（NVIDIA GPU的基本执行单元，一个warp包含32个线程）
#define WARP_SIZE 32
// 宏定义：将内存地址指向的4个连续int32元素转换为int4向量类型（CUDA的4元素整数向量）
// 用于SIMD优化，一次处理4个整数
#define INT4(value) (reinterpret_cast<int4 *>(&(value))[0])
// 宏定义：将内存地址指向的4个连续float元素转换为float4向量类型（CUDA的4元素浮点向量）
// 预留用于浮点型数据的SIMD处理
#define FLOAT4(value) (reinterpret_cast<float4 *>(&(value))[0])
这里#define INT4(value)把某个变量 reinterpret 为 int4 类型指针，然后取第一个 int4。
#define INT4(value) (reinterpret_cast<int4 *>(&(value))[0])
解释：int4 是 CUDA 提供的 4 个 int 的向量类型（4 × 32-bit = 128-bit）。reinterpret_cast<int4 *>(&(value))：将变量的地址视为 int4。[0]：访问这个向量。
用途：用于向量化加载/存储 4 个 int，一次操作 128-bit，提高内存吞吐量。比如：
int data[4] = {1,2,3,4};
int4 vec = INT4(data); // 将 data 当作 int4 使用
📉 2.3.3 int32 类型基础版直方图统计的 CUDA 核函数实现（histogram_i32_kernel）
这段代码是在 GPU 上统计整数出现次数的核心计算逻辑。简单说，就是让 GPU 的多个线程分工合作，快速数出每个数字出现了多少次。具体来说，给 GPU 分配了足够多的线程（每个线程块 256 个线程，总线程数根据数据量调整），每个线程负责处理一个数字。由于多个线程可能同时操作同一个位置（比如两个线程都遇到数字 1），这里用了 “原子操作” 保证计数准确（不会因为同时加 1 导致漏算）。最后，结果数组里每个位置的值就是对应数字出现的总次数。
功能：直方图统计的CUDA核函数（针对int32类型数据）
grid(N/256)：网格维度设置为输入数据量N除以256（每个网格包含若干线程块）block(256)：线程块维度为256（每个线程块包含256个线程）参数说明：
a: 输入数组（长度为N的int32类型GPU数组，存储待统计的整数，假设元素值>=1）y: 输出数组（直方图结果，存储每个整数的出现次数）N: 输入数组a的元素总数
// Histogram
// grid(N/256), block(256)
// a: Nx1, y: count histogram, a >= 1
__global__ void histogram_i32_kernel(int *a, int *y, int N) {
// 计算当前线程的全局索引：线程块索引×块内线程数 + 线程在块内的索引
// 每个线程负责处理输入数组中的一个元素
int idx = blockIdx.x * blockDim.x + threadIdx.x;
// 线程索引检查：仅当索引小于数据总量N时才执行（避免越界访问）
if (idx < N)
// 原子操作：将y数组中对应位置（a[idx]为索引）的值加1
// 原子操作确保多个线程同时更新同一位置时不会出现计数错误
atomicAdd(&(y[a[idx]]), 1);
}
🔢 2.3.4 基于向量计算（Vec4）优化的int32类型直方图统计CUDA核函数（histogram_i32x4_kernel）
这段代码是 GPU 直方图统计的 “提速版”，核心思路是让一个线程同时干 4 个活，比基础版效率更高。具体来说，基础版是 1 个线程处理 1 个数字，而这个优化版里，1 个线程一次处理 4 个连续数字。首先通过 “线程编号 ×4” 找到要处理的 4 个数字的起始位置，然后用向量操作把这 4 个数字一次性读到 GPU 的寄存器里（就像一次搬 4 个快递，比一次次单独搬快）。接着，线程分别给这 4 个数字对应的统计位置各加 1，并且用 “原子操作” 保证多个线程同时统计同一个数字时不会算错。为了不浪费 GPU 资源，线程总数也做了调整：基础版 1 个线程块 256 个线程，这个版本因为 1 个线程顶 4 个用，所以线程块里只放 64 个线程，总处理能力和基础版匹配，但内存访问和计算效率提升了不少。
基于向量计算（Vec4）优化的int32类型直方图统计CUDA核函数
grid(N/256)：网格维度设置为输入数据量N除以256（与基础版一致，控制总计算规模）block(256/4)：线程块维度为64（256除以4，因每个线程处理4个元素，总算力与基础版匹配）参数说明：
a: 输入数组（长度为N的int32类型GPU数组，存储待统计整数，假设元素值>=1）y: 输出数组（直方图结果，存储每个整数的出现次数）N: 输入数组a的元素总数
// Histogram + Vec4
// grid(N/256), block(256/4)
// a: Nx1, y: count histogram, a >= 1
__global__ void histogram_i32x4_kernel(int *a, int *y, int N) {
// 计算当前线程的“基索引”：每个线程处理4个连续元素，故索引需×4
// 公式含义：线程块索引×块内线程数（64） + 线程块内索引 → 得到线程处理的“4元素组编号”，再×4得到首个元素的全局索引
int idx = 4 * (blockIdx.x * blockDim.x + threadIdx.x);
// 索引检查：确保当前线程处理的首个元素不越界（避免访问超出N范围的数据）
if (idx < N) {
// 向量加载：将a数组中从idx开始的4个连续int32元素，一次性加载为CUDA的int4向量类型（reg_a为寄存器变量）
// 利用向量操作减少内存访问次数，提升效率（基础版需4次单独加载，此处1次完成）
int4 reg_a = INT4(a[idx]);
// 对向量的4个分量（x、y、z、w）分别执行原子加1操作
// 每个分量对应一个待统计的整数，通过原子操作确保计数不冲突，实现“1个线程统计4个元素”
atomicAdd(&(y[reg_a.x]), 1);
atomicAdd(&(y[reg_a.y]), 1);
atomicAdd(&(y[reg_a.z]), 1);
atomicAdd(&(y[reg_a.w]), 1);
}
}
为什么idx要*4？ FLOAT4 会把连续的 4 个元素当作一个向量处理。乘 4 是为了确保每个线程负责的 4 个元素不重叠。
#define FLOAT4(value) (reinterpret_cast<float4 *>(&(value))[0])
举例说明（thread=0，1，2）：
| thread | 处理元素 | x4后idx | 处理元素 |
|---|---|---|---|
| 0 | a[0:3] | 0 | a[0:3] |
| 1 | a[1:4] | 4 | a[4:7] |
| 2 | a[2:5] | 8 | a[8:11] |
🦭 2.3.5 把 CUDA kernel 封装成 PyTorch 可调用的扩展函数
剩余代码主要是将不同版本的 CUDA histogram kernel封装成 PyTorch 扩展函数，并通过 PyBind11 注册到 Python 中。
核心代码主要是以下内容：
- 用
TORCH_BINDING_ELEM_ADD自动生成一系列 C++ 封装函数。- 用
PYBIND11_MODULE和m.def把这些函数绑定到 Python模块里。- 最终，Python 可以直接调用这些扩展函数，而内部会执行高性能 CUDA 内核。
// 宏定义1：将输入的标识符转换为字符串（例如把func转为"func"）
// 用于后续绑定函数时，统一函数名的字符串表示
#define STRINGFY(str) #str
// 宏定义2：通用的PyTorch函数绑定宏
// 功能：将C++函数func绑定到PyTorch模块m中，函数名（字符串形式）与函数本身同名
// 作用：简化多个函数的重复绑定代码，避免手动写重复的def调用
#define TORCH_BINDING_COMMON_EXTENSION(func) \
m.def(STRINGFY(func), &func, STRINGFY(func));
// 宏定义3：PyTorch张量数据类型检查宏
// 功能：检查张量T的数据类型是否等于指定的th_type
// 若不匹配：打印张量当前的类型信息，并抛出运行时错误（提示需使用th_type类型）
#define CHECK_TORCH_TENSOR_DTYPE(T, th_type) \
if (((T).options().dtype() != (th_type))) { \
std::cout << "Tensor Info:" << (T).options() << std::endl; \
throw std::runtime_error("values must be " #th_type); \
}
// 宏定义4：PyTorch张量形状（第0维度）检查宏
// 功能：检查张量T第0维度的大小是否等于指定的S0
// 若不匹配：抛出运行时错误（提示张量大小不匹配）
#define CHECK_TORCH_TENSOR_SHAPE(T, S0) \
if (((T).size(0) != (S0))) { \
throw std::runtime_error("Tensor size mismatch!"); \
}
这段代码的核心作用是 “搭桥梁”—— 把之前写的 GPU 直方图统计核函数，变成 Python 里能直接调用的函数，还顺便做了数据检查和参数配置，保证调用时不出错。具体逻辑分三步：首先，用一堆宏定义简化重复工作，比如检查输入数据类型对不对、自动生成不同版本的统计函数；然后，针对基础版（1 个线程处理 1 个数据）和优化版（1 个线程处理 4 个数据），分别生成对应的 C++ 函数 —— 这些函数会先检查输入数据是否符合要求（比如必须是 int32 类型），再根据输入数据的最大值确定直方图的长度，接着配置 GPU 线程的分工（多少个线程块、每个块多少线程），最后调用对应的 CUDA 核函数去计算；最后，把这两个 C++ 函数绑定到 PyTorch 模块里，这样 Python 里导入模块后，就能像调用普通函数一样用 GPU 算直方图了。整个过程就像给 GPU 核函数套了个 “安全易用的外壳”，既保证了计算正确，又让用户不用管底层的 GPU 配置细节。
// 宏定义5：直方图统计函数的通用生成宏（核心宏）
// 参数说明：
// packed_type：函数名后缀（如i32、i32x4，对应不同实现）
// th_type：输入张量要求的数据类型（如torch::kInt32）
// element_type：数据的基础类型（如int，用于指针转换）
// n_elements：每个线程处理的元素数量（1对应基础版，4对应SIMD优化版）
#define TORCH_BINDING_HIST(packed_type, th_type, element_type, n_elements) \
// 定义直方图统计的C++函数（函数名：histogram_+packed_type，如histogram_i32）
torch::Tensor histogram_##packed_type(torch::Tensor a) { \
// 第一步：检查输入张量a的数据类型是否符合要求（如必须是int32）
CHECK_TORCH_TENSOR_DTYPE(a, (th_type)) \
// 第二步：配置输出张量y的参数（int32类型、与输入a同设备的GPU）
auto options = \
torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA, 0); \
// 第三步：获取输入张量a的元素总数N
const int N = a.size(0); \
// 第四步：计算输入张量a中的最大值（用于确定输出直方图的长度）
std::tuple<torch::Tensor, torch::Tensor> max_a = torch::max(a, 0); \
// 第五步：初始化输出张量y（长度为M+1，初始值0，存储直方图结果）
torch::Tensor max_val = std::get<0>(max_a).cpu(); \
const int M = max_val.item().to<int>(); \
auto y = torch::zeros({M + 1}, options); \
// 第六步：计算CUDA核函数的线程块大小（256除以每个线程处理的元素数，匹配算力）
static const int NUM_THREADS_PER_BLOCK = 256 / (n_elements); \
// 第七步：计算CUDA核函数的网格大小（向上取整，确保覆盖所有元素）
const int NUM_BLOCKS = (N + 256 - 1) / 256; \
// 第八步：定义CUDA的线程块和网格维度
dim3 block(NUM_THREADS_PER_BLOCK); \
dim3 grid(NUM_BLOCKS); \
// 第九步：调用对应的CUDA核函数（如histogram_i32_kernel）
// 传入输入a、输出y的GPU内存指针，以及元素总数N
histogram_##packed_type##_kernel<<<grid, block>>>( \
reinterpret_cast<element_type *>(a.data_ptr()), \
reinterpret_cast<element_type *>(y.data_ptr()), N); \
// 第十步：返回最终的直方图结果张量y
return y; \
}
最后代码用宏定义快速生成了两个 C++ 接口函数，分别对接基础版和 SIMD 优化版的 CUDA 直方图核函数，还自动做了数据类型检查、GPU 线程配置；最后通过模块绑定，让 Python 能直接调用这两个函数，不用管底层 CUDA 细节，轻松实现 GPU 加速的直方图统计。
// 调用宏生成两个具体的直方图函数：
// 1. histogram_i32：基础版，每个线程处理1个int32元素
// 2. histogram_i32x4：SIMD优化版，每个线程处理4个int32元素
TORCH_BINDING_HIST(i32, torch::kInt32, int, 1)
TORCH_BINDING_HIST(i32x4, torch::kInt32, int, 4)
// PyTorch扩展模块定义（模块名由TORCH_EXTENSION_NAME指定，即编译时的"hist_lib"）
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
// 将生成的两个直方图函数（histogram_i32、histogram_i32x4）绑定到模块m
// 绑定后，Python中可通过hist_lib.histogram_i32调用
TORCH_BINDING_COMMON_EXTENSION(histogram_i32)
TORCH_BINDING_COMMON_EXTENSION(histogram_i32x4)
}
⛪️ 三、小结
⭐️⭐️⭐️ 路漫漫其修远兮，吾将上下而求索。Fighting！⭐️⭐️⭐️
扫一扫
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
