[ET-VK] Modify quantized linear naive shader to linearly dispatch work to improve performance.

Pull Request resolved: #10116

This diff changes naive quantized linear mat mul op to use push constant instead of uniform buffers and change dispatch pattern to linear to improve performance.
ghstack-source-id: 278469019
@exported-using-ghexport

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

Author: trivedivivek

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

@@ -29,16 +29,20 @@ ${layout_declare_tensor(2, "r", "t_qmat2", "int8", STORAGE)}
 ${layout_declare_tensor(3, "r", "t_scales", DTYPE, STORAGE)}
 
 $if STORAGE == "buffer":
-  ${layout_declare_ubo(4, "ivec4", "out_sizes")}
-  ${layout_declare_ubo(5, "ivec4", "out_strides")}
-  ${layout_declare_ubo(6, "int", "out_numel")}
-  ${layout_declare_ubo(7, "ivec4", "mat1_sizes")}
-  ${layout_declare_ubo(8, "ivec4", "mat1_strides")}
-  ${layout_declare_ubo(9, "ivec4", "qmat2_strides")}
-  ${layout_declare_ubo(10, "ivec4", "scales_strides")}
+  layout(push_constant) uniform restrict Block {
+    ivec4 out_sizes;
+    ivec4 out_strides;
+    ivec4 mat1_sizes;
+    ivec4 mat1_strides;
+    ivec4 qmat2_strides;
+    ivec4 scales_strides;
+    int out_numel;
+  };
 $else:
-  ${layout_declare_ubo(4, "ivec3", "out_limits")}
-  ${layout_declare_ubo(5, "ivec4", "mat1_sizes")}
+  layout(push_constant) uniform restrict Block {
+    ivec3 out_limits;
+    ivec4 mat1_sizes;
+  };
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -83,42 +87,40 @@ void main() {
 
 #else // USING_TEXTURE
 
-#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
-
 void main() {
-  const u16vec2 out_pos = u16vec2(
-    gl_GlobalInvocationID.x,
-    gl_GlobalInvocationID.y);
+  const ivec2 out_pos = ivec2(
+    gl_GlobalInvocationID.x % out_limits.x,
+    gl_GlobalInvocationID.x / out_limits.x);
 
-  if (out_pos.x >= out_limits.x || out_pos.y >= out_limits.y) {
+  if (out_pos.y >= out_limits.y) {
     return;
   }
 
-  const uint16_t qmat2_pos_x = out_pos.x;
+  const int 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));
+  const VEC4_T scales = load_texel(t_scales,  ivec3(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++)
+    int i = 0, x = 0;
+    i < mat1_sizes.x;
+    i += 4, x++)
   {
-    mat1_tex = load_texel(t_mat1, u16vec3(x, out_pos.y, 0));
+    mat1_tex = load_texel(t_mat1, ivec3(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));
+    mat2_tex[0] = load_texel(t_qmat2, ivec3(out_pos.x, i, 0));
+    mat2_tex[1] = load_texel(t_qmat2, ivec3(out_pos.x, i + 1, 0));
+    mat2_tex[2] = load_texel(t_qmat2, ivec3(out_pos.x, i + 2, 0));
+    mat2_tex[3] = load_texel(t_qmat2, ivec3(out_pos.x, i + 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;
-  write_texel(t_out, u16vec3(out_pos, 0), outtex);
+  write_texel(t_out, ivec3(out_pos, 0), outtex);
 }
 
 #endif

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

@@ -98,47 +98,22 @@ void add_q_8w_linear_node(
   add_dtype_suffix(kernel_name, graph.dtype_of(out_W_packed));
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out_W_packed));
 
-  vkapi::ParamsBindList ubos({});
+  std::vector<PushConstantDataInfo> pcs;
   if (graph.is_buffer_storage(out_W_packed)) {
-    ubos.append(
-        {graph.sizes_ubo(out_W_packed),
-         graph.strides_ubo(out_W_packed),
-         graph.numel_ubo(out_W_packed),
-         graph.sizes_ubo(mat1_W_packed),
-         graph.strides_ubo(mat1),
-         graph.strides_ubo(q_mat2),
-         graph.strides_ubo(scales)});
+    pcs = {graph.sizes_pc_of(out_W_packed),
+      graph.strides_pc_of(out_W_packed),
+      graph.sizes_pc_of(mat1_W_packed),
+      graph.strides_pc_of(mat1),
+      graph.strides_pc_of(q_mat2),
+      graph.strides_pc_of(scales),
+      graph.numel_pc_of(out_W_packed)};
   } else {
-    ubos.append(
-        {graph.logical_limits_ubo(out_W_packed),
-         graph.sizes_ubo(mat1_W_packed)});
+    pcs = {graph.logical_limits_pc_of(out_W_packed),
+         graph.sizes_pc_of(mat1_W_packed)};
   }
 
-  utils::uvec3 global_wg;
-  if (graph.is_buffer_storage(out)) {
-    global_wg = {static_cast<uint32_t>(graph.numel_of(out_W_packed)), 1, 1};
-  } else {
-    global_wg = graph.logical_limits_of(out_W_packed);
-  }
-
-  utils::uvec3 local_wg{8, 8, 1};
-  int32_t out_W = graph.size_at<int32_t>(-1, out_W_packed);
-
-  if (graph.is_buffer_storage(out_W_packed)) {
-    local_wg[0] = 64;
-    local_wg[1] = 1;
-    local_wg[2] = 1;
-  } else {
-    if (out_W % 8 != 0) {
-      if (out_W % 4 == 0) {
-        local_wg[0] = 4;
-        local_wg[1] = 16;
-      } else {
-        local_wg[0] = 2;
-        local_wg[1] = 32;
-      }
-    }
-  }
+  const utils::uvec3 global_wg = {static_cast<uint32_t>(graph.numel_of(out_W_packed)), 1, 1};
+  const utils::uvec3 local_wg{64, 1, 1};
 
   graph.execute_nodes().emplace_back(new DispatchNode(
       graph,
@@ -149,11 +124,13 @@ void add_q_8w_linear_node(
       {{out_W_packed, vkapi::MemoryAccessType::WRITE},
        {{mat1_W_packed, q_mat2, scales}, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
-      ubos,
+      {},
       // Specialization Constants
       {},
       // Resizing Logic
-      resize_q_8w_linear_node));
+      resize_q_8w_linear_node,
+      {},
+      pcs));
   if (!graph.is_buffer_storage(out) &&
       graph.packed_dim_of(out) != WHCN::kWidthDim) {
     viewFn(graph, {out_W_packed, graph.add_none(), out});