forked from ccf-ai-infra/Intro-ops
63 lines
1.8 KiB
Plaintext
63 lines
1.8 KiB
Plaintext
#pragma once
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <float.h>
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
|
|
namespace oprt::softmax::nvidia {
|
|
|
|
__global__ void softmax_rowwise_kernel(float *out, const float *in,
|
|
int64_t rows, int64_t cols) {
|
|
// TODO: implement a numerically stable row-wise softmax kernel.
|
|
//
|
|
// Suggested steps:
|
|
// 1. Use one block per row.
|
|
// 2. Reduce to find the row maximum.
|
|
// 3. Compute exp(x - row_max), write the temporary values to out,
|
|
// and accumulate their sum.
|
|
// 4. Reduce to get the row sum.
|
|
// 5. Normalize each output element by row_sum.
|
|
int row = blockIdx.x;
|
|
extern __shared__ float smem[];
|
|
float thread_max = -INFINITY;
|
|
for (int col = threadIdx.x; col < cols; col += blockDim.x) {
|
|
if (in[row * cols + col] > thread_max) {
|
|
thread_max = in[row * cols + col];
|
|
}
|
|
}
|
|
|
|
smem[threadIdx.x] = thread_max;
|
|
|
|
__syncthreads();
|
|
for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
|
|
if (threadIdx.x < offset) {
|
|
smem[threadIdx.x] = smem[threadIdx.x] > smem[threadIdx.x + offset]
|
|
? smem[threadIdx.x]
|
|
: smem[threadIdx.x + offset];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
float row_max = smem[0];
|
|
float sum = 0.0f;
|
|
for (int col = threadIdx.x; col < cols; col += blockDim.x) {
|
|
out[row * cols + col] = exp(in[row * cols + col] - row_max);
|
|
sum += out[row * cols + col];
|
|
}
|
|
|
|
smem[threadIdx.x] = sum;
|
|
__syncthreads();
|
|
for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
|
|
if (threadIdx.x < offset) {
|
|
smem[threadIdx.x] += smem[threadIdx.x + offset];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
for (int col = threadIdx.x; col < cols; col += blockDim.x) {
|
|
out[row * cols + col] /= smem[0];
|
|
}
|
|
}
|
|
|
|
} // namespace oprt::softmax::nvidia
|