Clean up max_pool2d (#2575)

Summary:
Pull Request resolved: #2575

In anticipation of convolution, I noticed some nits in the max_pool2d implementation:
1. Removed unneeded shader branch.
2. Renamed `kernel` to `kernel_size`.
3. Removed `Int` to `IntList` normalization since this is already performed by the compiler, so the serialized graph will always have `IntList`. The fact that the Python tests still pass when they specify `Int` for `padding=0` and `dilation=1` affirms this. Renamed the helper function appropriately and improve readability.
ghstack-source-id: 219657438
bypass-github-export-checks

Reviewed By: copyrightly

Differential Revision: D55220203

fbshipit-source-id: f4d163e13826a7abb9ec07aaa4ef3c574826e948

Author: jorgep31415

backends/vulkan/runtime/graph/ops/ExecuteNode.h

@@ -77,7 +77,6 @@ class ExecuteNode final {
   const api::utils::uvec3 global_workgroup_size_;
   const api::utils::uvec3 local_workgroup_size_;
   const std::vector<ArgGroup> args_;
-  // TODO(T180906457): allow re-computing param buffers.
   std::vector<std::shared_ptr<api::UniformParamsBuffer>> params_;
   const ResizeFunction resize_fn_;
   const std::vector<ValueRef> resize_args_;

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

@@ -30,7 +30,7 @@ layout(set = 0, binding = 4) uniform PRECISION restrict InExtents {
 in_extents;
 
 layout(set = 0, binding = 5) uniform PRECISION restrict Params {
-  ivec2 kernel;
+  ivec2 kernel_size;
   ivec2 stride;
   ivec2 padding;
   ivec2 dilation;
@@ -49,7 +49,7 @@ void main() {
   const ivec2 ipos = pos.xy * params.stride - params.padding;
 
   const ivec2 start = ipos;
-  const ivec2 end = ipos + params.kernel * params.dilation;
+  const ivec2 end = ipos + params.kernel_size * params.dilation;
 
   vec4 out_texel = vec4(FLT_MIN);
   ivec4 idx_texel = ivec4(0);
@@ -66,9 +66,6 @@ void main() {
 
         out_texel = max(cur_texel, out_texel);
       }
-      else {
-        out_texel = max(vec4(FLT_MIN), out_texel);
-      }
     }
   }
 

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

@@ -37,24 +37,24 @@ void resize_max_pool2d_node(
   // Channel
   new_out_sizes.at(ndim - 3) = self.sizes().at(ndim - 3);
 
-  const auto kernel = normalize_wh(graph->get_val(extra_args[0]));
-  const auto stride = normalize_wh(graph->get_val(extra_args[1]));
-  const auto padding = normalize_wh(graph->get_val(extra_args[2]));
-  const auto dilation = normalize_wh(graph->get_val(extra_args[3]));
+  const auto kernel_size = reverse(*graph, extra_args[0]);
+  const auto stride = reverse(*graph, extra_args[1]);
+  const auto padding = reverse(*graph, extra_args[2]);
+  const auto dilation = reverse(*graph, extra_args[3]);
   const bool ceil_mode = graph->get_val(extra_args[4]).toBool();
 
   // Height
   new_out_sizes.at(ndim - 2) = calc_out_size(
       self.sizes().at(ndim - 2),
-      kernel.data[1],
+      kernel_size.data[1],
       stride.data[1],
       padding.data[1],
       dilation.data[1],
       ceil_mode);
   // Width
   new_out_sizes.at(ndim - 1) = calc_out_size(
       self.sizes().at(ndim - 1),
-      kernel.data[0],
+      kernel_size.data[0],
       stride.data[0],
       padding.data[0],
       dilation.data[0],
@@ -77,7 +77,7 @@ void check_max_pool2d_args(const vTensor& in, const vTensor& out) {
 void add_max_pool2d_node(
     ComputeGraph& graph,
     const ValueRef in,
-    const ValueRef kernel,
+    const ValueRef kernel_size,
     const ValueRef stride,
     const ValueRef padding,
     const ValueRef dilation,
@@ -99,10 +99,10 @@ void add_max_pool2d_node(
   apply_dtype_suffix(kernel_name, t_out);
 
   KernelParams kernel_params{
-      normalize_wh(graph.get_val(kernel)),
-      normalize_wh(graph.get_val(stride)),
-      normalize_wh(graph.get_val(padding)),
-      normalize_wh(graph.get_val(dilation)),
+      reverse(graph, kernel_size),
+      reverse(graph, stride),
+      reverse(graph, padding),
+      reverse(graph, dilation),
   };
 
   graph.execute_nodes().emplace_back(new ExecuteNode(
@@ -121,7 +121,7 @@ void add_max_pool2d_node(
       },
       // Resizing
       resize_max_pool2d_node,
-      {kernel, stride, padding, dilation, ceil_mode}));
+      {kernel_size, stride, padding, dilation, ceil_mode}));
 }
 
 void max_pool2d(ComputeGraph& graph, const std::vector<ValueRef>& args) {

backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.cpp

@@ -28,13 +28,9 @@ int64_t calc_out_size(
   return out_size;
 }
 
-api::utils::ivec2 normalize_wh(Value& v) {
-  if (v.isInt()) {
-    return api::utils::make_ivec2({v.toInt(), v.toInt()});
-  } else {
-    auto l = v.toIntList();
-    return api::utils::make_ivec2({l.at(1), l.at(0)});
-  }
+api::utils::ivec2 reverse(ComputeGraph& graph, ValueRef vref) {
+  return api::utils::make_ivec2(
+      graph.get_val(vref).toIntList(), /*reverse=*/true);
 }
 
 } // namespace vulkan

backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.h

@@ -12,28 +12,30 @@
 
 #include <ATen/native/vulkan/api/api.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
 
 namespace at {
 namespace native {
 namespace vulkan {
 
 struct KernelParams final {
-  api::utils::ivec2 kernel;
+  api::utils::ivec2 kernel_size;
   api::utils::ivec2 stride;
   api::utils::ivec2 padding;
   api::utils::ivec2 dilation;
 };
 
 int64_t calc_out_size(
     const int64_t in_size,
-    const int64_t kernel,
+    const int64_t kernel_size,
     const int64_t stride,
     const int64_t padding,
     const int64_t dilation,
     const bool ceil_mode);
 
-api::utils::ivec2 normalize_wh(Value& v);
+api::utils::ivec2 reverse(ComputeGraph& graph, ValueRef vref);
 
 } // namespace vulkan
 } // namespace native