Reduce conv2d_pw global wg size (#4045)

Summary:
Pull Request resolved: #4045

Since our conv2d_pw shader uses 2x2 tiles over the output, we can reduce the number of shader invocations, i.e., reduce the global workgroup size.
ghstack-source-id: 231375285

Reviewed By: SS-JIA

Differential Revision: D58956301

fbshipit-source-id: 34510c4584b5cf343d6a0c62d3fe42d27ec934e6

Author: jorgep31415

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

@@ -66,10 +66,10 @@ void main() {
   // | pos[2] | pos[3] |
   // +--------+--------+
   ivec3 pos[${TILE_SIZE * TILE_SIZE}];
-  for (int y = 0, i = 0; y < 2; ++y) {
-    for (int x = 0; x < 2; ++x) {
+  for (int y = 0, i = 0; y < ${TILE_SIZE}; ++y) {
+    for (int x = 0; x < ${TILE_SIZE}; ++x) {
       pos[i] = ivec3(
-          gpos.x * 2 + x, gpos.y * ${TILE_SIZE} + y, gpos.z);
+          gpos.x * ${TILE_SIZE} + x, gpos.y * ${TILE_SIZE} + y, gpos.z);
       i++;
     }
   }

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

@@ -289,6 +289,21 @@ Conv2dMethod get_conv2d_method(
   return Conv2dMethod::SlidingWindow;
 }
 
+api::utils::uvec3 create_conv2d_global_wg_size(
+    ComputeGraph& graph,
+    const Conv2dMethod method,
+    const ValueRef out) {
+  if (method == Conv2dMethod::Pointwise) {
+    const api::utils::uvec3 image_extents = graph.image_extents_of(out);
+    return {
+        api::utils::div_up(image_extents.data[0u], 2u),
+        api::utils::div_up(image_extents.data[1u], 2u),
+        image_extents.data[2u]};
+  } else {
+    return graph.create_global_wg_size(out);
+  }
+}
+
 void add_conv2d_node(
     ComputeGraph& graph,
     const ValueRef in,
@@ -357,7 +372,7 @@ void add_conv2d_node(
   graph.execute_nodes().emplace_back(new ExecuteNode(
       graph,
       shader,
-      graph.create_global_wg_size(out),
+      create_conv2d_global_wg_size(graph, method, out),
       graph.create_local_wg_size(out),
       // Inputs and Outputs
       {{out, api::MemoryAccessType::WRITE},

backends/vulkan/test/op_tests/cases.py

@@ -312,6 +312,17 @@ def get_conv_inputs():
                 [0],
                 5,
             ),
+            (
+                (1, 16, 672, 512),
+                (64, 16, 1, 1),
+                (64,),
+                [1, 1],
+                [0, 0],
+                [1, 1],
+                False,
+                [0, 0],
+                1,
+            ),
         ]
     )
     return test_suite