Increase the limit of buffer_size for matmul

Author: junjihashimoto

experimental/kernels/kernels.h

@@ -282,17 +282,31 @@ fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
     let T : u32 = params.T;
     let C : u32 = params.C;
     let OC : u32 = params.OC;
+    // N == B*T == global_id.x
     let b : u32 = global_id.x / T;
     let t : u32 = global_id.x % T;
-    if (b < B && t < T) {
-        let bt : u32 = b * T + t;
-        for (var o : u32 = 0u; o < OC; o++) {
-            var val : {{precision}} = bias[o];
-            for (var i : u32 = 0u; i < C; i++) {
-                val += inp[bt * C + i] * weight[o * C + i];
-            }
-            out[bt * OC + o] = val;
-        }
+    if (arrayLength(&bias) == 1) {
+      if (b < B && t < T) {
+          let bt : u32 = global_id.x;
+          for (var o : u32 = 0u; o < OC; o++) {
+              var val : {{precision}} = 0;
+              for (var i : u32 = 0u; i < C; i++) {
+                  val += inp[bt * C + i] * weight[o * C + i];
+              }
+              out[bt * OC + o] = val;
+          }
+      }
+    } else {
+      if (b < B && t < T) {
+          let bt : u32 = global_id.x;
+          for (var o : u32 = 0u; o < OC; o++) {
+              var val : {{precision}} = bias[o];
+              for (var i : u32 = 0u; i < C; i++) {
+                  val += inp[bt * C + i] * weight[o * C + i];
+              }
+              out[bt * OC + o] = val;
+          }
+      }
     }
 }
 )";

experimental/kernels/kernels_c.cpp

@@ -171,15 +171,57 @@ void MATMUL_FORWARD_GPU(float* out,
   unsigned long c = static_cast<unsigned long>(C);
   unsigned long oc = static_cast<unsigned long>(OC);
   setLogLevel(kError);
-  Context ctx = createContext();
-  Tensor inp_t = createTensor(ctx, Shape{b * t * c}, kf32, inp);
-  Tensor weight_t = createTensor(ctx, Shape{oc * c}, kf32, weight);
-  Tensor bias_t = createTensor(ctx, Shape{oc}, kf32, bias);
-  Tensor out_t = createTensor(ctx, Shape{b * t * oc}, kf32);
+  // See https://github.com/google/dawn/blob/a8fbe981a86cb59536e2de423d2013a82d9b54a0/src/dawn/native/Limits.cpp
+  WGPURequiredLimits requiredLimits = {
+    .limits = {
+      .maxTextureDimension1D=8192,
+      .maxTextureDimension2D=8192,
+      .maxTextureDimension3D=2048,
+      .maxTextureArrayLayers=256,
+      .maxBindGroups=4,
+      .maxBindGroupsPlusVertexBuffers=24,
+      .maxBindingsPerBindGroup=1000,
+      .maxDynamicUniformBuffersPerPipelineLayout=8,
+      .maxDynamicStorageBuffersPerPipelineLayout=4,
+      .maxSampledTexturesPerShaderStage=16,
+      .maxSamplersPerShaderStage=16,
+      .maxStorageBuffersPerShaderStage=8,
+      .maxStorageTexturesPerShaderStage=4,
+      .maxUniformBuffersPerShaderStage=12,
+      .maxUniformBufferBindingSize=65536,
+      .maxStorageBufferBindingSize=1073741824,
+      .minUniformBufferOffsetAlignment=256,
+      .minStorageBufferOffsetAlignment=256,
+      .maxVertexBuffers=8,
+      .maxBufferSize=0x80000000,
+      .maxVertexAttributes=16,
+      .maxVertexBufferArrayStride=2048,
+      .maxInterStageShaderComponents=64,
+      .maxInterStageShaderVariables=16,
+      .maxColorAttachments=8,
+      .maxColorAttachmentBytesPerSample=32,
+      .maxComputeWorkgroupStorageSize=16384,
+      .maxComputeInvocationsPerWorkgroup=256,
+      .maxComputeWorkgroupSizeX=256,
+      .maxComputeWorkgroupSizeY=256,
+      .maxComputeWorkgroupSizeZ=64,
+      .maxComputeWorkgroupsPerDimension=65535
+    },
+    .nextInChain = nullptr
+  };
+  Context ctx = createContext({},{},{
+      .requiredLimits = &requiredLimits
+    });
+
+  Tensor inp_i = createTensor(ctx, Shape{b * t * c}, kf32, inp);
+  Tensor weight_i = createTensor(ctx, Shape{oc * c}, kf32, weight);
+  Tensor bias_i = bias == NULL ? createTensor(ctx, Shape{1}, kf32) : createTensor(ctx, Shape{oc}, kf32, bias);
+  Tensor out_o = createTensor(ctx, Shape{b * t * oc}, kf32);
   std::promise<void> promise;
   std::future<void> future = promise.get_future();
+  assert ( (b*t) % 256 == 0 );
   Kernel op = createKernel(ctx, {kShaderMatmul, 256, kf32},
-                           Bindings{inp_t, weight_t, bias_t, out_t},
+                           Bindings{inp_i, weight_i, bias_i, out_o},
                            /* nWorkgroups */ {cdiv(b * t, 256), 1, 1},
                            /* params */
                            MatmulParams{
@@ -190,7 +232,7 @@ void MATMUL_FORWARD_GPU(float* out,
                            });
   dispatchKernel(ctx, op, promise);
   wait(ctx, future);
-  toCPU(ctx, out_t, out, b * t * oc * sizeof(float));
+  toCPU(ctx, out_o, out, b * t * oc * sizeof(float));
 }
 
 void MATMUL_BACKWARD_GPU(float* dinp, float* dweight, float* dbias,

experimental/kernels/kernels_c.h

@@ -10,7 +10,7 @@ extern "C" {
 #define USE_GPU_FOR_LAYERNORM_FORWARD 1
 // --  Note: atomicAdd should be used with i32 or u32 not f32.
 // #define USE_GPU_FOR_LAYERNORM_BACKWARD 1
-// #define USE_GPU_FOR_MATMUL_FORWARD 1
+#define USE_GPU_FOR_MATMUL_FORWARD 1
 // #define USE_GPU_FOR_MATMUL_BACKWARD 1
 #define USE_GPU_FOR_ATTENTION_FORWARD 1
 // #define USE_GPU_FOR_ATTENTION_BACKWARD 1