Update on "[ET-VK] Efficient tiled int8 matmul"

## Context

Introduce a optimized tiled implementation for computing the weight int8-quantized linear operation.

This implementation takes advantage of the following principles to squeeze out performance:

* Compute an output tile with each thread, rather than a single output element. This allows for better memory re-use of loaded input tensor data.
* Compute the output tile by iteratively loading tiles of the input matrices, caching them in registers, and then performing the `fma` accumulations to obtain a partial output. By splitting the data loading and computation into distinct steps, the GPU is able to perform latency hiding more effectively, i.e. switching to a warp that needs to perform compute when the current warp is waiting on data load
* Use a work group size of `{N, 1, 1}`. This makes it so that all the threads in a work group load the same row of the input matrx, and consecutive columns of the weight matrix. This way, the row of the input is kept hot in the cache, and accesses to the weight matrix can be coalesced due to the previous diff un-transposing the weight matrix.

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

[ghstack-poisoned]

Author: SS-JIA

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

@@ -47,13 +47,8 @@ ivec4 read_texel(ivec4 tidx) {
   ivec4 sizes_to_use = sizes;
   int packed_dim_to_use = packed_dim;
   if (transpose_hw == 1) {
-    int tmp = sizes_to_use.x;
-    sizes_to_use.x = sizes_to_use.y;
-    sizes_to_use.y = tmp;
-
-    tmp = tidx_to_use.x;
-    tidx_to_use.x = tidx.y;
-    tidx_to_use.y = tmp;
+    sizes_to_use.xy = sizes_to_use.yx;
+    tidx_to_use.xy = tidx.yx;
 
     if (packed_dim == 1) {
       packed_dim_to_use = 0;

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

@@ -33,13 +33,8 @@ void main() {
 
   ivec4 sizes = out_sizes;
   if (transpose_hw == 1) {
-    int tmp = sizes.x;
-    sizes.x = sizes.y;
-    sizes.y = tmp;
-
-    tmp = out_tidx.x;
-    out_tidx.x = out_tidx.y;
-    out_tidx.y = tmp;
+    sizes.xy = sizes.yx;
+    out_tidx.xy = out_tidx.yx;
   }
   const int in_nchwi = tidx_to_nchwi(out_tidx, sizes);
 

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

@@ -40,13 +40,8 @@ VEC4_T read_texel(ivec4 tidx) {
   ivec4 sizes_to_use = sizes;
   int packed_dim_to_use = packed_dim;
   if (transpose_hw == 1) {
-    int tmp = sizes_to_use.x;
-    sizes_to_use.x = sizes_to_use.y;
-    sizes_to_use.y = tmp;
-
-    tmp = tidx_to_use.x;
-    tidx_to_use.x = tidx.y;
-    tidx_to_use.y = tmp;
+    sizes_to_use.xy = sizes_to_use.yx;
+    tidx_to_use.xy = tidx.yx;
 
     if (packed_dim == 1) {
       packed_dim_to_use = 0;

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

@@ -70,15 +70,15 @@ void main() {
   // TODO(ssjia): optimize memory access pattern by traversing mat1 x in inner loop
   for (int i = 0; i < mat1_sizes.x; i++) {
     const FLOAT_T mat1_val = t_mat1[mat1_offset];
-    const FLOAT_T mat2_val = FLOAT_T(t_qmat2[qmat2_offset]) * scale;
+    const FLOAT_T mat2_val = FLOAT_T(t_qmat2[qmat2_offset]);
 
     outval += mat1_val * mat2_val;
 
     mat1_offset++;
     qmat2_offset += qmat2_strides.y;
   }
 
-  t_out[out_bufi] = outval;
+  t_out[out_bufi] = outval * scale;
 }
 
 #else // USING_TEXTURE

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

@@ -175,8 +175,8 @@ ValueRef prepack_standard_hw_transposed(
   const int w_dim = new_out_sizes.size() - 1;
   const int h_dim = new_out_sizes.size() - 2;
   const int64_t tmp = new_out_sizes.at(w_dim);
-  new_out_sizes[w_dim] = new_out_sizes[h_dim];
-  new_out_sizes[h_dim] = tmp;
+  new_out_sizes.at(w_dim) = new_out_sizes.at(h_dim);
+  new_out_sizes.at(h_dim) = tmp;
   ValueRef tensor = graph.add_tensor(
       new_out_sizes,
       graph.dtype_of(tensor_data),