38 lines
1.0 KiB
Plaintext
38 lines
1.0 KiB
Plaintext
#ifndef __SOFTMAX_KERNEL_CUH__
|
|
#define __SOFTMAX_KERNEL_CUH__
|
|
#define WINDOW 129
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
#include <float.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
|
|
__global__ void local_softmax_forward_kernel(float* out, float inv_temperature, const float* inp, int N, int T) {
|
|
// Implement this
|
|
int x = blockDim.x * blockIdx.x + threadIdx.x;
|
|
|
|
if(x < N*T){
|
|
//row offset based on WIndow size
|
|
int rowoff = x * WINDOW;
|
|
|
|
float maxval = -FLT_MAX;
|
|
for (int i = 0; i < WINDOW; i++) {
|
|
maxval = fmaxf(maxval, inp[i + rowoff]);
|
|
}
|
|
|
|
float sumval = 0.0f;
|
|
//Traverse Window Length Again
|
|
for (int i = 0; i < WINDOW; i++) {
|
|
float ev = expf(inv_temperature * (inp[i + rowoff] - maxval));
|
|
sumval += ev;
|
|
out[i + rowoff] = ev;
|
|
}
|
|
|
|
float norm = 1.0f / sumval;
|
|
for (int i = 0; i < WINDOW; i++) {
|
|
out[i + rowoff] *= norm;
|
|
}
|
|
}
|
|
}
|
|
|
|
#endif // __SOFTMAX_KERNEL_CUH__ |