94 lines
2.8 KiB
Plaintext
94 lines
2.8 KiB
Plaintext
// Softmax and log-softmax CUDA kernels -- f32 native.
|
|
// Common header (common_device_functions.cuh) is prepended by build.rs
|
|
// providing: float, (), expf(), logf(), fmaxf(), etc.
|
|
//
|
|
// Each warp/block handles one row of [dim] elements.
|
|
|
|
extern "C" __global__ void softmax_forward(
|
|
float* __restrict__ output,
|
|
const float* __restrict__ input,
|
|
int batch_size,
|
|
int dim
|
|
) {
|
|
int row = blockIdx.x;
|
|
if (row >= batch_size) return;
|
|
|
|
const float* in_row = input + row * dim;
|
|
float* out_row = output + row * dim;
|
|
|
|
// Find max for numerical stability
|
|
float max_val = (-1e4f);
|
|
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
|
float v = in_row[i];
|
|
max_val = fmaxf(max_val, v);
|
|
}
|
|
// Warp reduction for max
|
|
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
|
|
float other = __shfl_down_sync(0xffffffff, max_val, offset);
|
|
max_val = fmaxf(max_val, other);
|
|
}
|
|
// Broadcast max from lane 0
|
|
max_val = (__shfl_sync(0xffffffff, (max_val), 0));
|
|
|
|
// Compute exp(x - max) and sum
|
|
float sum = 0.0f;
|
|
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
|
float e = expf(in_row[i] - max_val);
|
|
out_row[i] = e;
|
|
sum = sum + e;
|
|
}
|
|
// Warp reduction for sum
|
|
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
|
|
sum = sum + __shfl_down_sync(0xffffffff, sum, offset);
|
|
}
|
|
sum = (__shfl_sync(0xffffffff, (sum), 0));
|
|
|
|
// Normalize
|
|
float inv_sum = 1.0f / sum;
|
|
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
|
out_row[i] = out_row[i] * inv_sum;
|
|
}
|
|
}
|
|
|
|
extern "C" __global__ void log_softmax_forward(
|
|
float* __restrict__ output,
|
|
const float* __restrict__ input,
|
|
int batch_size,
|
|
int dim
|
|
) {
|
|
int row = blockIdx.x;
|
|
if (row >= batch_size) return;
|
|
|
|
const float* in_row = input + row * dim;
|
|
float* out_row = output + row * dim;
|
|
|
|
// Find max for numerical stability
|
|
float max_val = (-1e4f);
|
|
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
|
float v = in_row[i];
|
|
max_val = fmaxf(max_val, v);
|
|
}
|
|
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
|
|
float other = __shfl_down_sync(0xffffffff, max_val, offset);
|
|
max_val = fmaxf(max_val, other);
|
|
}
|
|
max_val = (__shfl_sync(0xffffffff, (max_val), 0));
|
|
|
|
// Compute sum of exp(x - max)
|
|
float sum = 0.0f;
|
|
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
|
sum = sum + expf(in_row[i] - max_val);
|
|
}
|
|
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
|
|
sum = sum + __shfl_down_sync(0xffffffff, sum, offset);
|
|
}
|
|
sum = (__shfl_sync(0xffffffff, (sum), 0));
|
|
|
|
float log_sum = logf(sum);
|
|
|
|
// log_softmax(x_i) = x_i - max - log(sum(exp(x - max)))
|
|
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
|
out_row[i] = in_row[i] - max_val - log_sum;
|
|
}
|
|
}
|