Softmax¶
row_max¶
每个线程只访问数组的一个数据，所以循环每轮只依次比较整个数组的前后半部分数据：
for(int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) {
out[threadIdx.x] = fmaxf(out[threadIdx.x], out[threadIdx.x + stride]);
}
__syncthreads();
}
cols可能大于blockDim.x，所以设置内部循环；rows可能大于gridDim.x，所以设置外层循环。
多个 block 同时处理多个 row；每个 block 内，多个线程同时处理该 row 的多个 col。
线程先分别计算自己负责的列的最大值，得到每个线程的最大值。然后通过block_reduce_max进一步计算block的最大值，也即行最大值:
for (int row = blockIdx.x; row < rows; row += gridDim.x) {
const float* input_row = matrix + static_cast<size_t>(row) * cols;
float* output_row = out + static_cast<size_t>(row) * cols;
float thread_max = -INFINITY;
for (int col = threadIdx.x; col < cols; col += blockDim.x) {
thread_max = fmaxf(thread_max, input_row[col]);
}
const float row_max =
block_reduce_max(thread_max, warp_results);
warp 内归约。每个warp 32个线程，利用shuffle指令：
for (int offset = 16; offset > 0; offset >>= 1) {
value = fmaxf(
value,
__shfl_down_sync(0xffffffff, value, offset));
}
在shuffle指令开始前，warp会保证32个线程都已经完成局部最大值的计算，所以不需要使用syncthreads().
warp归约的结果在lane 0线程中：
warp 归约 只有warp 0参与此归约过程。 如果有256个线程，则共有8个warp参与block_reduce_max，把8个warp的归约结果取出，用warp 0进行再一次归约，得到block最大值。
if (warp_id == 0){
value = (lane < num_warps) ? warp_results[lane] : 0.0f;
value = warp_reduce_max(value);
if (lane == 0){
warp_results[0] = value;
}
}
完整示例代码：
#include<cstdio>
#include<cstdlib>
#include<cmath>
#include<cuda_runtime.h>
#include<algorithm>
#include<vector>
#define CUDA_CHECK(call) do { \
cudaError_t err = (call); \
if (err != cudaSuccess) { \
std::fprintf(stderr, "%s:%d CUDA error: %s\n", \
__FILE__, __LINE__, cudaGetErrorString(err)); \
std::exit(1); \
} \
} while (0)
__device__ __forceinline__
float warp_reduce_max(float value) {
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
value = fmaxf(
value,
__shfl_down_sync(0xffffffff, value, offset));
}
return value;
}
__device__ __forceinline__
float warp_reduce_sum(float value) {
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
value += __shfl_down_sync(0xffffffff, value, offset);
}
return value;
}
__device__ __forceinline__
float block_reduce_max(float value, float* warp_results) {
const int lane = threadIdx.x & 31;
const int warp_id = threadIdx.x >> 5;
const int num_warps = (blockDim.x + 31) >> 5;
value = warp_reduce_max(value);
if (lane == 0) {
warp_results[warp_id] = value;
}
__syncthreads();
if (warp_id == 0) {
value = lane < num_warps ? warp_results[lane] : -INFINITY;
value = warp_reduce_max(value);
if (lane == 0) {
warp_results[0] = value;
}
}
__syncthreads();
return warp_results[0];
}
__device__ __forceinline__
float block_reduce_sum(float value, float* warp_results) {
const int lane = threadIdx.x & 31;
const int warp_id = threadIdx.x >> 5;
const int num_warps = (blockDim.x + 31) >> 5;
value = warp_reduce_sum(value);
if (lane == 0) {
warp_results[warp_id] = value;
}
__syncthreads();
if (warp_id == 0) {
value = lane < num_warps ? warp_results[lane] : 0.0f;
value = warp_reduce_sum(value);
if (lane == 0) {
warp_results[0] = value;
}
}
__syncthreads();
return warp_results[0];
}
__global__ void softmax_v2(
const float* __restrict__ matrix,
float* __restrict__ out,
int rows,
int cols) {
// 最多支持 1024 线程，即最多 32 个 warp
__shared__ float warp_results[32];
for (int row = blockIdx.x; row < rows; row += gridDim.x) {
const float* input_row = matrix + static_cast<size_t>(row) * cols;
float* output_row = out + static_cast<size_t>(row) * cols;
float thread_max = -INFINITY;
for (int col = threadIdx.x; col < cols; col += blockDim.x) {
thread_max = fmaxf(thread_max, input_row[col]);
}
const float row_max =
block_reduce_max(thread_max, warp_results);
float thread_sum = 0.0f;
for (int col = threadIdx.x; col < cols; col += blockDim.x) {
// 更重视精度时换成 expf
const float value = __expf(input_row[col] - row_max);
output_row[col] = value;
thread_sum += value;
}
const float row_sum =
block_reduce_sum(thread_sum, warp_results);
const float inverse_sum = 1.0f / row_sum;
for (int col = threadIdx.x; col < cols; col += blockDim.x) {
output_row[col] *= inverse_sum;
}
// grid-stride 处理下一行前，确保所有线程结束当前行。
__syncthreads();
}
}
void softmax_cpu(
const std::vector<float>& input,
std::vector<float>& output,
int rows,
int cols
) {
for (int row = 0; row < rows; ++row) {
const float* input_row = input.data() + row * cols;
float* output_row = output.data() + row * cols;
float max_value = input_row[0];
for (int col = 1; col < cols; ++col) {
max_value = std::max(max_value, input_row[col]);
}
double sum = 0.0;
for (int col = 0; col < cols; ++col) {
output_row[col] =
static_cast<float>(std::exp(input_row[col] - max_value));
sum += output_row[col];
}
for (int col = 0; col < cols; ++col) {
output_row[col] /= static_cast<float>(sum);
}
}
}
int main(){
const int rows = 3;
const int cols = 5;
const size_t bytes = static_cast<size_t>(rows) * cols * sizeof(float);
std::vector<float> input = {
1.0f, 2.0f, 3.0f, 4.0f, 5.0f,
-2.0f, -1.0f, 0.0f, 1.0f, 2.0f,
1000.0f, 1001.0f, 999.0f, 998.0f, 997.0f
};
std::vector<float> gpu_output(rows * cols);
std::vector<float> cpu_output(rows * cols);
float* d_input = nullptr;
float* d_output = nullptr;
CUDA_CHECK(cudaMalloc(&d_input, bytes));
CUDA_CHECK(cudaMalloc(&d_output, bytes));
CUDA_CHECK(cudaMemcpy(
d_input, input.data(), bytes, cudaMemcpyHostToDevice
));
CUDA_CHECK(cudaMemset(d_output, 0, bytes));
// Reduction requires a power-of-two block with at least one thread per col.
int threads = 1;
while(threads < cols){
threads <<= 1;
}
dim3 grid(rows);
dim3 block(threads);
size_t shared_bytes = static_cast<size_t>(threads) * sizeof(float);
softmax_v2<<<grid, block, shared_bytes>>>(d_input, d_output, rows, cols);
// 启动错误
CUDA_CHECK(cudaGetLastError());
// 执行期间的错误
CUDA_CHECK(cudaDeviceSynchronize());
CUDA_CHECK(cudaMemcpy(
gpu_output.data(),
d_output,
bytes,
cudaMemcpyDeviceToHost
));
softmax_cpu(input, cpu_output, rows, cols);
bool passed = true;
constexpr float atol = 1e-5f;
constexpr float rtol = 1e-4f;
for (int i = 0; i < rows * cols; ++i) {
float error = std::fabs(gpu_output[i] - cpu_output[i]);
float tolerance = atol + rtol * std::fabs(cpu_output[i]);
if (!std::isfinite(gpu_output[i]) || error > tolerance) {
std::printf(
"Mismatch at index %d: GPU=%g, CPU=%g, error=%g\n",
i, gpu_output[i], cpu_output[i], error
);
passed = false;
}
}
// Softmax 每一行的和应该约等于 1
for (int row = 0; row < rows; ++row) {
float sum = 0.0f;
for (int col = 0; col < cols; ++col) {
sum += gpu_output[row * cols + col];
}
std::printf("row %d sum = %.8f\n", row, sum);
if (std::fabs(sum - 1.0f) > 1e-5f) {
passed = false;
}
}
std::printf("%s\n", passed ? "TEST PASSED" : "TEST FAILED");
CUDA_CHECK(cudaFree(d_input));
CUDA_CHECK(cudaFree(d_output));
return passed ? EXIT_SUCCESS : EXIT_FAILURE;
}