Update on "[ET-VK] Add coop shader for int8 linear"

Title says it all!

## Changes

* Apply co-operative shader for vector * matrix computations.

Differential Revision: [D73279548](https://our.internmc.facebook.com/intern/diff/D73279548/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/runtime/graph/ops/glsl/q_8w_linear_coop.glsl

@@ -57,15 +57,15 @@ void main() {
   VEC4_T b[4];
   VEC4_T local_c[TILE_ROWS];
 
+  [[unroll]] for (int i = 0; i < TILE_ROWS; ++i) {
+    local_c[i] = VEC4_T(0.0);
+  }
+
   $if SCALES_STORAGE == "buffer":
     const VEC4_T scales = VEC4_T(t_scales[out_col >> 2]);
   $else:
     const VEC4_T scales = VEC4_T(texelFetch(t_scales, ivec2(out_col >> 2, 0), 0));
 
-  [[unroll]] for (int i = 0; i < TILE_ROWS; ++i) {
-    partial_c[gid][wid][i] = VEC4_T(0.0);
-  }
-
   for (int pos = 4 * wid; pos < in_sizes.x; pos += (4 * NWORKERS)) {
     // Preload t_weight
     [[unroll]] for (int i = 0; i < 4; i++) {
@@ -77,17 +77,17 @@ void main() {
     // Preload t_in
     for (int i = 0; i < TILE_ROWS; i++) {
       $if IN_STORAGE == "buffer":
-        a[i] = t_in[((out_row + i) * in_sizes.x + ((pos)) >> 2)];
+        a[i] = t_in[((out_row + i) * in_sizes.x + pos) >> 2];
       $else:
         a[i] = VEC4_T(texelFetch(t_in, ivec3(pos >> 2, out_row + i, 0), 0));
     }
 
-    // Compute t_out...?
+    // Accumulate partial output
     [[unroll]] for (int i = 0; i < TILE_ROWS; ++i) {
-        local_c[i] += a[i].x * b[0]
-                    + a[i].y * b[1]
-                    + a[i].z * b[2]
-                    + a[i].w * b[3];
+        local_c[i] += a[i].x * b[0] +
+                      a[i].y * b[1] +
+                      a[i].z * b[2] +
+                      a[i].w * b[3];
     }
   }
 

backends/vulkan/runtime/graph/ops/glsl/q_8w_linear_tiled.glsl

@@ -71,18 +71,18 @@ void main() {
     // Preload input tensor
     [[unroll]] for (int i = 0; i < TILE_ROWS; i++) {
       $if IN_STORAGE == "buffer":
-        a[i] = t_in[((out_row + i) * in_sizes.x + (pos)) >> 2];
+        a[i] = t_in[((out_row + i) * in_sizes.x + pos) >> 2];
       $else:
         a[i] = VEC4_T(texelFetch(t_in, ivec3(pos >> 2, out_row + i, 0), 0));
     }
 
-    // Compute partial output
+    // Accumulate output
     [[unroll]] for (int i = 0; i < TILE_ROWS; ++i) {
         c[i] += a[i].x * b[0] + a[i].y * b[1] + a[i].z * b[2] + a[i].w * b[3];
     }
   }
 
-  // Store output tensor
+  // Store to output tensor
   [[unroll]] for (int i = 0; i < TILE_ROWS; ++i) {
     $if OUT_STORAGE == "buffer":
       if (out_row + i < out_sizes.y) {

backends/vulkan/test/op_tests/cases.py

@@ -153,25 +153,24 @@ def get_linear_inputs():
 def get_weight_int8pack_mm_inputs():
     MKN_list = [
         [1, 480, 256],
-        # [1, 1024, 1024],
-        # [1, 1024, 256],
-        # [3, 480, 256],
-        # [6, 480, 256],
-        # [6, 256, 1024],
-        # [6, 1024, 256],
-        # [6, 256, 256],
-        # [6, 256, 512],
-        # [4, 768, 4096],
-        # [1024, 1024, 1024],
+        [1, 1024, 1024],
+        [1, 1024, 256],
+        [3, 480, 256],
+        [6, 480, 256],
+        [6, 256, 1024],
+        [6, 1024, 256],
+        [6, 256, 256],
+        [6, 256, 512],
+        [4, 768, 4096],
+        [1024, 1024, 1024],
     ]
 
     inputs_list = [((M, K), (N, K), (N)) for M, K, N in MKN_list]
 
     test_suite = VkTestSuite(inputs_list)
     test_suite.dtypes = ["at::kFloat"]
     test_suite.layouts = ["utils::kWidthPacked"]
-    # test_suite.storage_types = ["utils::kTexture3D", "utils::kBuffer"]
-    test_suite.storage_types = ["utils::kBuffer"]
+    test_suite.storage_types = ["utils::kTexture3D", "utils::kBuffer"]
     test_suite.prepacked_args = ["mat2", "scales"]
     test_suite.requires_prepack = True