[ET-VK] Modify quantized linear tiling shader to linearly dispatch work to improve thread occupancy and performance.

Pull Request resolved: #10508

This diff changes tiled 8 bit quantized linear mat mul op to linearly dispatch work which increases thread occupancy and improves performance.

Differential Revision: [D73751979](https://our.internmc.facebook.com/intern/diff/D73751979/)
ghstack-source-id: 280651869

Author: trivedivivek

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

@@ -36,13 +36,16 @@ layout(push_constant) uniform restrict Block {
   ivec4 weight_sizes;
 };
 
+#include "indexing_utils.h"
+
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 void main() {
-  const uint out_row = gl_GlobalInvocationID.y * TILE_ROWS;
-  const uint out_col = gl_GlobalInvocationID.x << 2;
+  const uint out_size_x_div_4 = divup4(out_sizes.x);
+  const uint out_col = (gl_GlobalInvocationID.x % out_size_x_div_4) << 2;
+  const uint out_row = (gl_GlobalInvocationID.x / out_size_x_div_4) * TILE_ROWS;
 
-  if (out_col >= out_sizes.x || out_row >= out_sizes.y) {
+  if (out_row >= out_sizes.y) {
     return;
   }
 

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

@@ -16,10 +16,10 @@ q_8w_linear_tiled:
     TILE_ROWS:
       - VALUE: 1
         SUFFIX: o4x1
+      - VALUE: 2
+        SUFFIX: o4x2
       - VALUE: 4
         SUFFIX: o4x4
-      - VALUE: 6
-        SUFFIX: o4x6
   shader_variants:
     - NAME: q_8w_linear_tiled_texture3d_texture3d_texture2d_texture2d_float
     - NAME: q_8w_linear_tiled_buffer_buffer_texture2d_texture2d_float

backends/vulkan/runtime/graph/ops/impl/QuantizedLinearInt8.cpp

@@ -182,8 +182,8 @@ void add_q_8w_linear_tiled_node(
   const int64_t M = utils::val_at(-2, mat1_sizes);
   int out_tile_nrows = 4;
   if (M % 6 == 0) {
-    kernel_name += "_o4x6";
-    out_tile_nrows = 6;
+    kernel_name += "_o4x2";
+    out_tile_nrows = 2;
   } else if (M % 4 == 0) {
     kernel_name += "_o4x4";
     out_tile_nrows = 4;
@@ -197,6 +197,10 @@ void add_q_8w_linear_tiled_node(
 
   utils::uvec3 global_wg_size = graph.logical_limits_of(out);
   global_wg_size[1] = global_wg_size[1] / out_tile_nrows;
+  if (!use_coop_algorithm) {
+    global_wg_size[0] *= global_wg_size[1];
+    global_wg_size[1] = 1;
+  }
 
   utils::uvec3 local_wg_size{64, 1, 1};
   if (use_coop_algorithm) {