Disable chromium.subgroup_matrix_uniformity

Author: junjihashimoto

examples/matmul/run.cpp

@@ -618,6 +618,7 @@ inline KernelCode createMatmulWithTranspose(const char *shaderTemplate, const si
 inline KernelCode createMatmul12(const char *shaderTemplate, const size_t M,
                                  const size_t K, const size_t N,
                                  const size_t TM, const size_t TN,
+                                 const size_t LID,
                                  const Shape &workgroupSize = {256, 1, 1},
                                  NumType precision = kf32) {
   std::string codeString(shaderTemplate);
@@ -626,7 +627,8 @@ inline KernelCode createMatmul12(const char *shaderTemplate, const size_t M,
                           {"{{K}}", toString(K)},
                           {"{{N}}", toString(N)},
                           {"{{TM}}", toString(TM)},
-                          {"{{TN}}", toString(TN)}
+                          {"{{TN}}", toString(TN)},
+                          {"{{LID}}", toString(LID)}
     });
   return {loopUnrolling(codeString), workgroupSize, precision};
 }
@@ -638,18 +640,18 @@ inline KernelCode createMatmul12(const char *shaderTemplate, const size_t M,
 const char* kShaderSubgroupMatrixMultiply = R"(
 enable subgroups;
 enable chromium_experimental_subgroup_matrix;
+diagnostic (off, chromium.subgroup_matrix_uniformity);
 
 @group(0) @binding(0) var<storage, read_write>  A: array<{{precision}}>;
 @group(0) @binding(1) var<storage, read_write>  B: array<{{precision}}>;
 @group(0) @binding(2) var<storage, read_write>  C: array<{{precision}}>;
 
 @compute @workgroup_size({{workgroupSize}})
-fn main(@builtin(workgroup_id) wg: vec3<u32>) {
+fn main(@builtin(workgroup_id) wg: vec3<u32>,
+        @builtin(local_invocation_id) localID : vec3<u32>) {
 
   let rowStart: u32 = wg.x * 8u * {{TM}};
-  let colStart: u32 = wg.y * 8u * {{TN}};
-
-  if (rowStart >= u32({{M}}) || colStart >= u32({{N}})) { return; }
+  let colStart: u32 = (wg.y * {{LID}} + localID.y)  * 8u * {{TN}};
 
   let baseA: u32 = rowStart * {{K}};
   let baseB: u32 = colStart;
@@ -661,27 +663,41 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>) {
   // 4x4 accumulators (8x8 each)
   var accxx: array<subgroup_matrix_result<{{precision}}, 8, 8>, {{TM}} * {{TN}}>;
 
+  for (var idx_i: u32 = 0; idx_i < {{TM}}; idx_i++) {
+    Ax[idx_i] = subgroup_matrix_left<{{precision}}, 8, 8>(0);
+  }
+
+  for (var idx_i: u32 = 0; idx_i < {{TN}}; idx_i++) {
+    Bx[idx_i] = subgroup_matrix_right<{{precision}}, 8, 8>(0);
+  }
+
+  for (var idx_i: u32 = 0; idx_i < {{TM}}; idx_i++) {
+    for (var idx_j: u32 = 0; idx_j < {{TN}}; idx_j++) {
+      accxx[idx_i+idx_j*{{TM}}] = subgroup_matrix_result<{{precision}}, 8, 8>(0);
+    }
+  }
+
   for (var k: u32 = 0u; k < {{K}}; k = k + 8u) {
     workgroupBarrier();
     for (var idx_i: u32 = 0; idx_i < {{TM}}; idx_i++) {
-      Ax[idx_i] = subgroupMatrixLoad<subgroup_matrix_left<{{precision}},8,8>>(&A, baseA +  idx_i * 8u*{{K}} + k, false, {{K}});
+      Ax[idx_i] = subgroupMatrixLoad<subgroup_matrix_left<{{precision}},8,8>>(&A, baseA + k + 8u * {{K}} * idx_i, false, {{K}});
     }
 
     for (var idx_i: u32 = 0; idx_i < {{TN}}; idx_i++) {
-      Bx[idx_i] = subgroupMatrixLoad<subgroup_matrix_right<{{precision}},8,8>>(&B, baseB + k*{{N}} +  8u * idx_i, false, {{N}});
+      Bx[idx_i] = subgroupMatrixLoad<subgroup_matrix_right<{{precision}},8,8>>(&B, baseB + k * {{N}} + 8u * idx_i, false, {{N}});
     }
 
-    for (var idx_i: u32 = 0; idx_i < {{TM}}; idx_i++) {
-      for (var idx_j: u32 = 0; idx_j < {{TN}}; idx_j++) {
-        accxx[idx_i+idx_j*{{TM}}] = subgroupMatrixMultiplyAccumulate(Ax[idx_i], Bx[idx_j], accxx[idx_i+idx_j*{{TM}}]);
+    for (var idx_j: u32 = 0; idx_j < {{TN}}; idx_j++) {
+      for (var idx_i: u32 = 0; idx_i < {{TM}}; idx_i++) {
+        accxx[idx_j*{{TM}} + idx_i] = subgroupMatrixMultiplyAccumulate(Ax[idx_i], Bx[idx_j], accxx[idx_j*{{TM}} + idx_i]);
       }
     }
   }
 
   workgroupBarrier();
   for (var idx_i: u32 = 0; idx_i < {{TM}}; idx_i++) {
     for (var idx_j: u32 = 0; idx_j < {{TN}}; idx_j++) {
-      subgroupMatrixStore(&C, cBase + idx_i * 8u * {{N}} + 8u * idx_j, accxx[idx_i+idx_j*{{TM}}], false, {{N}});
+      subgroupMatrixStore(&C, cBase + idx_i * 8u * {{N}} + 8u * idx_j, accxx[idx_j*{{TM}} + idx_i], false, {{N}});
     }
   }
 }
@@ -858,15 +874,16 @@ Kernel selectMatmul(Context &ctx, int version,
                           /*nWorkgroups*/ nWorkgroups,
                           NoParam{}, &info);
   } else if (version == 12) {
-    // f32: Subgroup matrix multiply
-    static constexpr size_t TM = 2;
-    static constexpr size_t TN = 4;
-    Shape wgSize = {64, 1, 1}; // One subgroup per workgroup
-    Shape nWorkgroups = {cdiv(M, 8 * TM), cdiv(N, 8 * TN), 1};
+    // f16: Subgroup matrix multiply
+    static constexpr size_t TM = 4;
+    static constexpr size_t TN = 8;
+    static constexpr size_t LID = 2;
+    Shape wgSize = {64, LID, 1}; // One subgroup per workgroup
+    Shape nWorkgroups = {cdiv(M, 8 * TM), cdiv(N, 8 * TN * LID), 1};
     LOG(kDefLog, kInfo, "M: %zu, K: %zu, N: %zu", M, K, N);
     LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str());
     LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str());
-    KernelCode matmul = createMatmul12(kShaderSubgroupMatrixMultiply, M, K, N, TM, TN, wgSize, numtype);
+    KernelCode matmul = createMatmul12(kShaderSubgroupMatrixMultiply, M, K, N, TM, TN, LID, wgSize, numtype);
     kernel = createKernel(ctx, matmul, bindings, nWorkgroups,
                           NoParam{}, &info);
   }
@@ -931,6 +948,10 @@ void runTest(int version, size_t M, size_t K, size_t N,
       LOG(kDefLog, kError, "[DeviceLost %d] %.*s\n", (int)reason, (int)msg.length, msg.data);
     }
   };
+
+  static WGPULimits requiredLimits = WGPU_LIMITS_INIT;
+
+  devDesc.requiredLimits = &requiredLimits;
 
   Context ctx = createContext({}, {}, devDesc);