33 lines
1.1 KiB
Plaintext
33 lines
1.1 KiB
Plaintext
#ifndef ENCODER_FORWARD_KERNEL_CUH
|
|
#define ENCODER_FORWARD_KERNEL_CUH
|
|
|
|
#include <cuda_runtime.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
|
|
__global__ void encoder_forward_kernel(float* __restrict__ out, const int* __restrict__ inp, const float* __restrict__ wte, const float* __restrict__ wpe,
|
|
int B, int T, int C) {
|
|
// Implement this
|
|
int x = blockDim.x * blockIdx.x + threadIdx.x;
|
|
int y = blockDim.y * blockIdx.y + threadIdx.y;
|
|
|
|
if(x < B && y < T){
|
|
int token = inp[x*T + y];
|
|
const float* __restrict__ wte_ptr = wte + token * C;
|
|
const float* __restrict__ wpe_ptr = wpe + y * C;
|
|
for(int i = 0; i < C; i ++){
|
|
out[x*T*C + y*C + i] = wte_ptr[i] + wpe_ptr[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
// Launch kernel here
|
|
void encoder_forward(float* out, const int* inp, const float* wte, const float* wpe, int B, int T, int C) {
|
|
// Implement this
|
|
dim3 block(16,16,1);
|
|
dim3 grid((B + block.x -1)/ block.x,(T+block.y-1)/block.y);
|
|
|
|
encoder_forward_kernel<<<grid, block>>>(out, inp, wte, wpe,B,T,C);
|
|
}
|
|
|
|
|
|
#endif // ENCODER_FORWARD_KERNEL_CUH |