[ET-VK] Removed shared memory usage and simplied conv2d dw op shader to improve performance. (#11270)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #11178 by
@trivedivivek
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/102/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/102/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/101/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/102/orig
@diff-train-skip-merge

---------

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

Author: pytorchbot

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

@@ -47,11 +47,6 @@ layout(push_constant) uniform restrict Block {
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
-// For performance improvement, reduce register usage by caching positions in shared memory.
-// Offset index by 1 every 16 points to avoid bank access conflict.
-#define offset_pos_index(index) (index + ((index) >> 4))
-shared ivec3 pos_shared[offset_pos_index(LOCAL_WG_SIZE)];
-
 /*
  * Computes a depthwise convolution. Each shader invocation calculates the
  * output at a single output location.
@@ -77,8 +72,6 @@ void main() {
     return;
   }
 
-  pos_shared[offset_pos_index(gl_LocalInvocationIndex)] = pos;
-
   // Compute the index of the top-left element of the overlay region. Negative
   // indices indicate that the top-left element is in a region added by padding.
   const ivec2 ipos = pos.xy * stride - padding;
@@ -89,13 +82,10 @@ void main() {
   const ivec2 end = ipos + overlay_region.xy;
 
   // sum outputs
-  VEC4_T sum[BATCH_SIZE_Y][BATCH_SIZE_X];
+  VEC4_T sum[BATCH_SIZE_Y * BATCH_SIZE_X];
 
-  sum[0][0] = texelFetch(t_bias, ivec2(pos.z, 0), 0);
-  for (int y = 0; y < BATCH_SIZE_Y; y++) {
-    for (int x = 0; x < BATCH_SIZE_X; x++) {
-      sum[y][x] = sum[0][0];
-    }
+  for (int i = 0; i < BATCH_SIZE_Y * BATCH_SIZE_X; i++) {
+    sum[i] = VEC4_T(0);
   }
 
   // array to store input texels
@@ -115,7 +105,7 @@ void main() {
     if (i > 0) {
       for (int j = 0; j < TILE_SIZE; j++) {
         for (int s = 0; s < BATCH_SIZE_X; s++) {
-          sum[1][s] = fma(in_texels[j + s], prev_kernel_line[j], sum[1][s]);
+          sum[BATCH_SIZE_X + s] = fma(in_texels[j + s], prev_kernel_line[j], sum[BATCH_SIZE_X + s]);
         }
       }
     }
@@ -125,19 +115,19 @@ void main() {
       for (int j = 0; j < TILE_SIZE; j++, kx++) {
         prev_kernel_line[j] = texelFetch(t_kernel, ivec2(kx, pos.z), 0);
         for (int s = 0; s < BATCH_SIZE_X; s++) {
-          sum[0][s] = fma(in_texels[j + s], prev_kernel_line[j], sum[0][s]);
+          sum[s] = fma(in_texels[j + s], prev_kernel_line[j], sum[s]);
         }
       }
     }
   }
 
-  const ivec3 out_pos = pos_shared[offset_pos_index(gl_LocalInvocationIndex)];
+  const VEC4_T bias = texelFetch(t_bias, ivec2(pos.z, 0), 0);
   for (int y = 0; y < BATCH_SIZE_Y; y++) {
     for (int x = 0; x < BATCH_SIZE_X; x++) {
-      if (any(greaterThanEqual(ivec3(out_pos.x + x, out_pos.y + y, out_pos.z), out_limits.xyz))) {
-        continue;
+      const ivec3 out_pos = ivec3(pos.x + x, pos.y + y, pos.z);
+      if (all(lessThan(out_pos.xy, out_limits.xy))) {
+        imageStore(t_out, out_pos, op(sum[y * BATCH_SIZE_X + x] + bias, out_min, out_max));
       }
-      imageStore(t_out, ivec3(out_pos.x + x, out_pos.y + y, out_pos.z), op(sum[y][x], out_min, out_max));
     }
   }
 }