[ET-VK] Minor tuning for conv2d pw op to improve performance.

Pull Request resolved: #11113

The diff introduces minor tuning for the Conv2d pointwise (PW) operation in the Vulkan backend to improve performance.
Conv 2d pw now issues a 2D dispatch instead of 1D, where dispatch axis y is now sized based on output texture's batch size.
ghstack-source-id: 286652099

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

Author: trivedivivek

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

@@ -49,12 +49,12 @@ void main() {
   const uint div_by_x = gl_GlobalInvocationID.x / out_limits_scaled.x;
   const ivec3 gpos = ivec3(
     gl_GlobalInvocationID.x % out_limits_scaled.x,
-    div_by_x % out_limits_scaled.y,
-    div_by_x / out_limits_scaled.y);
+    div_by_x,
+    gl_GlobalInvocationID.y);
 
   // If the top left position is out of bounds, then this invocation will have
   // no work to do.
-  if (gpos.z >= out_limits.z) {
+  if (gpos.y >= out_limits_scaled.y || gpos.z >= out_limits.z) {
     return;
   }
 

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

@@ -398,8 +398,10 @@ void add_conv2d_node(
   utils::uvec3 wg_size = create_conv2d_global_wg_size(
       graph, method, out, weight_data, stride_equals_dilation);
 
-  if (method == Conv2dMethod::Pointwise || method == Conv2dMethod::Depthwise) {
+  if (method == Conv2dMethod::Depthwise) {
     wg_size = {wg_size[0] * wg_size[1] * wg_size[2], 1, 1};
+  } else if (method == Conv2dMethod::Pointwise) {
+    wg_size = {wg_size[0] * wg_size[1], wg_size[2], 1};
   }
 
   vkapi::ParamsBindList param_buffers;