[ET-VK] Using shared variable to store calculated output pose to free up registers and improve performance. (#7503)

* [ET-VK] Reduced int precision for all int storage in conv pw op to improve performance.

Pull Request resolved: #7447

This diff reduces the precision of all int storage in the conv pw op to improve performance. The code changes include adding the extension GL_EXT_shader_explicit_arithmetic_types_int16 and changing the data type of ints to uint16.

ghstack-source-id: 260166244
@exported-using-ghexport

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

* [ET-VK] Minor fix to conv 2d op using wg_size from create_conv2d_global_wg_size to determine local wg size.

Pull Request resolved: #7450

This diff contains changes to the Convolution.cpp file in the Vulkan backend of Executorch. The changes involve updating the code to use the create_conv2d_global_wg_size function to determine the local workgroup size for the convolution operation. This is done to ensure that the correct workgroup size is used for the operation, which can improve performance.

ghstack-source-id: 260166246
@exported-using-ghexport

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

* [ET-VK] Modify conv 2d pw op shader and dispatch settings to linearly dispatch work accounting for linearity texture to improve performance.

Pull Request resolved: #7452

This diff modifies the convolution 2D pointwise op shader and dispatch settings to linearly dispatch work accounting for linearity texture to improve performance.
ghstack-source-id: 260166247
@exported-using-ghexport

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

* [ET-VK] Using vec2 to store output positions to reudce shader register footprint.

Pull Request resolved: #7474

The diff changes the use of `u16vec3` to `u16vec2` to store output positions in the conv2d_pw op. This change is made to reduce the shader register footprint and improve performance.
ghstack-source-id: 260166245
@exported-using-ghexport

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

* [ET-VK] Using shared variable to store calculated output pose to free up registers and improve performance.

Pull Request resolved: #7475

This diff introduces a shared variable to store calculated output pose in conv2d_pw op to free up registers and improve performance. The code changes include adding a shared variable to hold calculated positions and modifying the existing code to use the shared variable.
ghstack-source-id: 260166242

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

---------

Co-authored-by: Vivek Trivedi <[email protected]>

Author: pytorchbot

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

@@ -34,13 +34,17 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 
+// shared memory to hold calculated positions, this would reduce register usage thus improving performance.
+shared u16vec2 pos_shared[gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z * TILE_SIZE * TILE_SIZE];
+
 /*
  * Computes a 2D pointwise convolution of an NxN output tile. Calculating an
  * output tile for pointwise convolution is more efficient because the kernel
  * size is only 1x1, making it easier to re-use loaded texels from t_kernel.
  */
 void main() {
   const uint16_t out_limits_y_scaled = uint16_t((out_limits.y + TILE_SIZE - 1) / TILE_SIZE);
+  const uint shared_mem_stride = gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z;
 
   const u16vec3 gpos = u16vec3(
     gl_GlobalInvocationID.x / (out_limits_y_scaled * out_limits.z),
@@ -58,6 +62,7 @@ void main() {
     for (int x = 0; x < TILE_SIZE; ++x) {
       pos[i] = u16vec2(
           gpos.x * TILE_SIZE + x, gpos.y * TILE_SIZE + y);
+      pos_shared[(shared_mem_stride * i) + gl_LocalInvocationIndex] = pos[i];
       i++;
     }
   }
@@ -73,7 +78,7 @@ void main() {
   // the top-left element is in a region added by padding.
   u16vec2 ipos[TILE_SIZE * TILE_SIZE];
   for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) {
-    ipos[i] = pos[i].xy * u16vec2(stride) - u16vec2(padding);
+    ipos[i] = pos[i] * u16vec2(stride) - u16vec2(padding);
   }
 
   vec4 sum[TILE_SIZE * TILE_SIZE];
@@ -138,8 +143,9 @@ void main() {
   }
 
   for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) {
-    if (all(lessThan(u16vec3(pos[i], gpos.z), out_limits))) {
-      imageStore(t_out, u16vec3(pos[i], gpos.z), op(sum[i], out_min, out_max));
+    const u16vec2 pos = pos_shared[(shared_mem_stride * i) + gl_LocalInvocationIndex];
+    if (all(lessThan(u16vec3(pos, gpos.z), out_limits))) {
+      imageStore(t_out, u16vec3(pos, gpos.z), op(sum[i], out_min, out_max));
     }
   }
 }