Update on "[ET-VK] Store weights transposed for int8 linear"

## Context

The weight tensor of a linear layer is usually stored in a transposed manner, such that when computing the matrix multiplication, the reduction traverses along the rows of the weight tensor as opposed to the columns. This results in a better memory access pattern for CPUs.

However, for GPUs, I have found that "un-transposing" the weight tensors result in better performance. This is likely due to the fact since GPUs can compute multiple output elements in parallel, reading along the columns allows for coalescing memory loads among threads in a work group.

## Changes

* Introduce the ability to transpose height and weight dims when transferring tensor data to the GPU.
* Prepackthe weight tensor "un-transposed" for the int8 quantized linear operator

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

[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),