comp files

kernels/encoder.cuh

@@ -4,27 +4,31 @@
 #include <cuda_runtime.h>
 #include "../utils/cuda_utils.cuh"
 
-__global__ void encoder_forward_kernel(float* out, const int* inp, const float* wte, const float* wpe,
+__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];
-        for(int i = 0; i < C; i ++){
-            out[x*T*C + y*C + i] = wte[token*C+i] + wpe[y*C+i];
-        }
+    int row = blockIdx.x * blockDim.y + threadIdx.y;
+    if (row >= B * T){
+        return;
+    }
+    int b = row / T;
+    int t = row % T;
+    int token = inp[b * T + t];
+    const float* __restrict__ wte_ptr = wte + token * C;
+    const float* __restrict__ wpe_ptr = wpe + t * C;
+    float* __restrict__ out_ptr = out + row * C;
+    for (int i = threadIdx.x; i < C; i += 32) {
+        out_ptr[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);
+    dim3 block(32, 8, 1);
+    int total_rows = B * T;
+    dim3 grid((total_rows + 8 - 1) / 8, 1, 1);
+    encoder_forward_kernel<<<grid, block>>>(out, inp, wte, wpe, B, T, C);
 }