Merge branch 'main' into size_test_m

Author: digantdesai

backends/apple/coreml/scripts/install_requirements.sh

@@ -12,7 +12,7 @@ SCRIPT_DIR_PATH="$(
 
 # TODO(jathu): remove the need to fetch coremltools to build deps for coreml_executor_runner.
 # Keep this version in sync with: pyproject.toml
-COREMLTOOLS_VERSION="8.1"
+COREMLTOOLS_VERSION="8.2"
 
 red=`tput setaf 1`
 green=`tput setaf 2`

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

@@ -27,6 +27,8 @@ ${layout_declare_ubo(B, "ivec4", "sizes")}
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 ${layout_declare_spec_const(C, "int", "t_layout", "DEFAULT_LAYOUT")}
+${layout_declare_spec_const(C, "int", "transpose_hw", "0")}
+
 const lowp ivec4 axis_map = unhash_axis_map(t_layout);
 const lowp int packed_dim = unhash_packed_dim(t_layout);
 
@@ -41,8 +43,23 @@ int extend_sign(int x) {
 }
 
 ivec4 read_texel(ivec4 tidx) {
+  ivec4 tidx_to_use = tidx;
+  ivec4 sizes_to_use = sizes;
+  int packed_dim_to_use = packed_dim;
+  if (transpose_hw == 1) {
+    sizes_to_use.xy = sizes_to_use.yx;
+    tidx_to_use.xy = tidx.yx;
+
+    if (packed_dim == 1) {
+      packed_dim_to_use = 0;
+    }
+    if (packed_dim == 0) {
+      packed_dim_to_use = 1;
+    }
+  }
+
   const ivec4 buf_indices = tidx_to_nchwi(
-      tidx, sizes, packed_dim);
+      tidx_to_use, sizes_to_use, packed_dim_to_use);
 
   int shift = (1 << 8) - 1;
   ivec4 masks;
@@ -70,7 +87,7 @@ ivec4 read_texel(ivec4 tidx) {
 
 void main() {
   const ivec3 lpos = ivec3(gl_GlobalInvocationID);
-  const ivec4 tidx = lpos_to_tidx(lpos, sizes, axis_map.w, packed_dim);
+  ivec4 tidx = lpos_to_tidx(lpos, sizes, axis_map.w, packed_dim);
 
   if (any(greaterThanEqual(tidx, sizes))) {
     return;

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

@@ -21,6 +21,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 // This constant is unused in this shader but is kept so that the signature is
 // consistent with nchw_to_image.
 ${layout_declare_spec_const(C, "int", "UNUSED_layout", "0")}
+${layout_declare_spec_const(C, "int", "transpose_hw", "0")}
 
 void main() {
   int out_bufi = int(gl_GlobalInvocationID.x);
@@ -29,7 +30,13 @@ void main() {
   }
 
   ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides);
-  const int in_nchwi = tidx_to_nchwi(out_tidx, out_sizes);
+
+  ivec4 sizes = out_sizes;
+  if (transpose_hw == 1) {
+    sizes.xy = sizes.yx;
+    out_tidx.xy = out_tidx.yx;
+  }
+  const int in_nchwi = tidx_to_nchwi(out_tidx, sizes);
 
   t_out[out_bufi] = nchw_in[in_nchwi];
 }

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

@@ -30,14 +30,31 @@ $if not FROM_STAGING:
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 ${layout_declare_spec_const(C, "int", "t_layout", "DEFAULT_LAYOUT")}
+${layout_declare_spec_const(C, "int", "transpose_hw", "0")}
+
 const lowp ivec4 axis_map = unhash_axis_map(t_layout);
 const lowp int packed_dim = unhash_packed_dim(t_layout);
 
 VEC4_T read_texel(ivec4 tidx) {
+  ivec4 tidx_to_use = tidx;
+  ivec4 sizes_to_use = sizes;
+  int packed_dim_to_use = packed_dim;
+  if (transpose_hw == 1) {
+    sizes_to_use.xy = sizes_to_use.yx;
+    tidx_to_use.xy = tidx.yx;
+
+    if (packed_dim == 1) {
+      packed_dim_to_use = 0;
+    }
+    if (packed_dim == 0) {
+      packed_dim_to_use = 1;
+    }
+  }
+
   $if FROM_STAGING:
-    const ivec4 buf_indices = tidx_to_nchwi(tidx, sizes, packed_dim);
+    const ivec4 buf_indices = tidx_to_nchwi(tidx_to_use, sizes_to_use, packed_dim_to_use);
   $else:
-    const ivec4 buf_indices = tidx_to_4bufi(tidx, buf_strides, packed_dim);
+    const ivec4 buf_indices = tidx_to_4bufi(tidx_to_use, buf_strides, packed_dim_to_use);
 
   VEC4_T texel = VEC4_T(0);
   if (tidx[packed_dim] < sizes[packed_dim]) {

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

@@ -64,24 +64,21 @@ void main() {
 
   FLOAT_T outval = FLOAT_T(0.0);
 
-  // Initial mat1 tensor idx will be (0, out_tidx.y, out_tidx.z, 0)
   int mat1_offset = out_tidx.y * mat1_strides.y + out_tidx.z * qmat2_strides.z;
-  // Initial qmat2 tensor idx wil be (0, out_tidx.x, 0, 0); note that the qmat2
-  // tensor is transposed
-  int qmat2_offset = out_tidx.x * qmat2_strides.y;
+  int qmat2_offset = out_tidx.x;
 
   // 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 = 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_offset += qmat2_strides.y;
   }
 
-  t_out[out_bufi] = outval;
+  t_out[out_bufi] = outval * scale;
 }
 
 #else // USING_TEXTURE
@@ -97,25 +94,27 @@ void main() {
     return;
   }
 
-  const uint16_t qmat2_pos_y = out_pos.x * uint16_t(4);
+  const uint16_t qmat2_pos_x = out_pos.x;
 
   VEC4_T outtex = VEC4_T(0);
 
   const VEC4_T scales = load_texel(t_scales,  u16vec3(out_pos.x, 0, 0));
 
+  VEC4_T mat1_tex;
+  VEC4_T mat2_tex[4];
   for (
     uint16_t i = uint16_t(0), x = uint16_t(0);
     i < uint16_t(mat1_sizes.x);
     i += uint16_t(4), x++)
   {
-    const VEC4_T mat1_tex = load_texel(t_mat1, u16vec3(x, out_pos.y, 0));
-    const VEC4_T sums = VEC4_T(
-        dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y, 0))),
-        dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y + uint16_t(1), 0))),
-        dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y + uint16_t(2), 0))),
-        dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y + uint16_t(3), 0))));
-
-    outtex += sums;
+    mat1_tex = load_texel(t_mat1, u16vec3(x, out_pos.y, 0));
+
+    mat2_tex[0] = load_texel(t_qmat2, u16vec3(out_pos.x, i, 0));
+    mat2_tex[1] = load_texel(t_qmat2, u16vec3(out_pos.x, i + uint16_t(1), 0));
+    mat2_tex[2] = load_texel(t_qmat2, u16vec3(out_pos.x, i + uint16_t(2), 0));
+    mat2_tex[3] = load_texel(t_qmat2, u16vec3(out_pos.x, i + uint16_t(3), 0));
+
+    outtex += mat1_tex.x * mat2_tex[0] + mat1_tex.y * mat2_tex[1] + mat1_tex.z * mat2_tex[2] + mat1_tex.w * mat2_tex[3];
   }
 
   outtex *= scales;