Fix the data alignment of softmax_forward

Author: junjihashimoto

experimental/kernels/kernels.h

@@ -165,15 +165,17 @@ static const char *kShaderSoftmax1 = R"(
 struct Params {
     N: u32,
     C: u32,
+    Cp: u32,
 };
 const NEG_INFINITY: f32 = -3.0e38; // WGSL has problem representing -3.4028235e+38
 @compute @workgroup_size({{workgroupSize}})
 fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
     let N : u32 = params.N;
     let C : u32 = params.C;
+    let Cp : u32 = params.Cp;
     let i : u32 = global_id.x;
     if (i < N) {
-        let inp_row_start : u32 = i * C;
+        let inp_row_start : u32 = i * Cp;
         var maxval : f32 = NEG_INFINITY;
         // Find the maximum value in the row
         for (var j : u32 = 0u; j < C; j++) {
@@ -194,6 +196,9 @@ fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
         for (var j : u32 = 0u; j < C; j++) {
             out[inp_row_start + j] /= sum;
         }
+        for (var j : u32 = C; j < Cp; j++) {
+            out[inp_row_start + j] = 0;
+        }
     }
 }
 )";

experimental/kernels/kernels_c.cpp

@@ -379,25 +379,28 @@ void RESIDUAL_BACKWARD_GPU(float* dinp1, float* dinp2, float* dout, int N){
   toCPU(ctx, dinp2_o, dinp2, n * sizeof(float));
 }
 
-void SOFTMAX_FORWARD_GPU(float* probs, float* logits, int b, int t, int v, int vp) {
+void SOFTMAX_FORWARD_GPU(float* probs, float* logits, int B, int T, int V, int Vp) {
   struct SoftmaxParam {
     uint32_t N;
     uint32_t C;
+    uint32_t Cp;
   };
-  uint32_t B = static_cast<uint32_t>(b);
-  uint32_t T = static_cast<uint32_t>(t);
-  uint32_t C = static_cast<uint32_t>(v);
+  uint32_t b = static_cast<uint32_t>(B);
+  uint32_t t = static_cast<uint32_t>(T);
+  uint32_t c = static_cast<uint32_t>(V);
+  uint32_t cp = static_cast<uint32_t>(Vp);
   Context ctx = createContext();
-  Tensor input = createTensor(ctx, {B * T, C}, kf32, logits);
-  Tensor output = createTensor(ctx, {B * T, C}, kf32);
+  Tensor input = createTensor(ctx, {b * t, cp}, kf32, logits);
+  Tensor output = createTensor(ctx, {b * t, cp}, kf32);
   std::promise<void> promise;
   std::future<void> future = promise.get_future();
+  assert( (B*T) % 256 == 0);
   Kernel op = createKernel(
       ctx, {kShaderSoftmax1, 256, kf32}, Bindings{input, output},
-      Shape{cdiv(B * T, 256), 1, 1}, SoftmaxParam{B * T, C});
+      Shape{cdiv(B * T, 256), 1, 1}, SoftmaxParam{b * t, c, cp});
   dispatchKernel(ctx, op, promise);
   wait(ctx, future);
-  toCPU(ctx, output, probs, sizeof(float)*B*T*C);
+  toCPU(ctx, output, probs, sizeof(float)*b*t*cp);
 }
 
 void CROSSENTROPY_FORWARD_GPU(float* losses,

experimental/kernels/kernels_c.h

@@ -18,8 +18,7 @@ extern "C" {
 #define USE_GPU_FOR_GELU_BACKWARD 1
 #define USE_GPU_FOR_RESIDUAL_FORWARD 1
 #define USE_GPU_FOR_RESIDUAL_BACKWARD 1
-// --  Note: softmax_forward_gpu is implemented, but there are some mismatches when enabled. --
-// #define USE_GPU_FOR_SOFTMAX_FORWARD 1
+#define USE_GPU_FOR_SOFTMAX_FORWARD 1
 #define USE_GPU_FOR_CROSSENTROPY_FORWARD 1
 // #define USE_GPU_FOR_CROSSENTROPY_SOFTMAX_BACKWARD 1