Apply loop-unrolling

Author: junjihashimoto

examples/matmul/run.cpp

@@ -628,7 +628,7 @@ inline KernelCode createMatmul12(const char *shaderTemplate, const size_t M,
                           {"{{TM}}", toString(TM)},
                           {"{{TN}}", toString(TN)}
     });
-  return {codeString, workgroupSize, precision};
+  return {loopUnrolling(codeString), workgroupSize, precision};
 }
 
 // ─────────────────────────────────────────────────────────────────────────────
@@ -663,25 +663,25 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>) {
 
   for (var k: u32 = 0u; k < {{K}}; k = k + 8u) {
     workgroupBarrier();
-    for (var i: u32 = 0; i < {{TM}}; i++) {
-      Ax[i] = subgroupMatrixLoad<subgroup_matrix_left<{{precision}},8,8>>(&A, baseA +  i * 8u*{{K}} + k, false, {{K}});
+    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}});
     }
 
-    for (var i: u32 = 0; i < {{TN}}; i++) {
-      Bx[i] = subgroupMatrixLoad<subgroup_matrix_right<{{precision}},8,8>>(&B, baseB + k*{{N}} +  8u * i, false, {{N}});
+    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}});
     }
 
-    for (var i: u32 = 0; i < {{TM}}; i++) {
-      for (var j: u32 = 0; j < {{TN}}; j++) {
-        accxx[i+j*{{TM}}] = subgroupMatrixMultiplyAccumulate(Ax[i], Bx[j], accxx[i+j*{{TM}}]);
+    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}}]);
       }
     }
   }
 
   workgroupBarrier();
-  for (var i: u32 = 0; i < {{TM}}; i++) {
-    for (var j: u32 = 0; j < {{TN}}; j++) {
-      subgroupMatrixStore(&C, cBase + i * 8u * {{N}} + 8u * j, accxx[i+j*{{TM}}], false, {{N}});
+  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}});
     }
   }
 }