Update on "[ET-VK] Refactor Pool.cpp"

This change adds more lines than it subtracts, but it'll be worth it once we reuse the methods for `aten.convolution`.

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

[ghstack-poisoned]

Author: jorgep31415

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -123,13 +123,6 @@ ValueRef ComputeGraph::add_tensor(
   return idx;
 }
 
-ValueRef ComputeGraph::add_tensor(
-    TensorRef& tref,
-    const api::StorageType storage_type,
-    const api::GPUMemoryLayout memory_layout) {
-  return add_tensor(tref.sizes, tref.dtype, storage_type, memory_layout);
-}
-
 ValueRef ComputeGraph::add_tensor(
     const std::vector<int64_t>& sizes,
     const api::ScalarType dtype,
@@ -139,9 +132,18 @@ ValueRef ComputeGraph::add_tensor(
       sizes, dtype, suggested_storage_type(), memory_layout, shared_object_idx);
 }
 
-ValueRef ComputeGraph::add_tensor(
-    TensorRef& tref,
+ValueRef ComputeGraph::add_tensor_like(
+    const ValueRef vref,
+    const api::StorageType storage_type,
+    const api::GPUMemoryLayout memory_layout) {
+  TensorRef& tref = get_val(vref).toTensorRef();
+  return add_tensor(tref.sizes, tref.dtype, storage_type, memory_layout);
+}
+
+ValueRef ComputeGraph::add_tensor_like(
+    const ValueRef vref,
     const api::GPUMemoryLayout memory_layout) {
+  TensorRef& tref = get_val(vref).toTensorRef();
   return add_tensor(tref.sizes, tref.dtype, memory_layout);
 }
 

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -174,14 +174,6 @@ class ComputeGraph final {
       const api::GPUMemoryLayout memory_layout,
       const int64_t shared_object_idx = -1);
 
-  /*
-   * Add a `vTensor` value to the graph with the properties of `tref`.
-   */
-  ValueRef add_tensor(
-      TensorRef& tref,
-      const api::StorageType storage_type,
-      const api::GPUMemoryLayout memory_layout);
-
   /*
    * Add a `vTensor` value to the graph with the specified properties. The
    * suggested storage type will be used to construct the `vTensor`.
@@ -192,14 +184,6 @@ class ComputeGraph final {
       const api::GPUMemoryLayout memory_layout,
       const int64_t shared_object_idx = -1);
 
-  /*
-   * Add a `vTensor` value to the graph with the properties of `tref`. The
-   * suggested storage type will be used to construct the `vTensor`.
-   */
-  ValueRef add_tensor(
-      TensorRef& tref,
-      const api::GPUMemoryLayout memory_layout);
-
   /*
    * Add a `vTensor` value to the graph with the specified properties. The
    * suggested storage type and memory layout will be used to construct the
@@ -210,6 +194,22 @@ class ComputeGraph final {
       const api::ScalarType dtype,
       const int64_t shared_object_idx = -1);
 
+  /*
+   * Add a `vTensor` value to the graph with the properties of `vref`.
+   */
+  ValueRef add_tensor_like(
+      const ValueRef vref,
+      const api::StorageType storage_type,
+      const api::GPUMemoryLayout memory_layout);
+
+  /*
+   * Add a `vTensor` value to the graph with the properties of `vref`. The
+   * suggested storage type will be used to construct the `vTensor`.
+   */
+  ValueRef add_tensor_like(
+      const ValueRef vref,
+      const api::GPUMemoryLayout memory_layout);
+
   /*
    * Add a `TensorRef` value to the graph with the specific properties. A
    * `TensorRef` is a reference to a `vTensor` whose data is stored in an

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

@@ -35,16 +35,16 @@ void resize_max_pool2d_node(
   new_out_sizes.at(ndim - 3) = self.sizes().at(ndim - 3);
 
   // Height, Width
-  const auto hw_sizes = calc_hw_out_sizes(
+  const auto new_out_sizes_hw = calc_out_sizes_hw(
       *graph,
       self.sizes(),
       extra_args[0],
       extra_args[1],
       extra_args[2],
       extra_args[3],
       extra_args[4]);
-  new_out_sizes.at(ndim - 2) = hw_sizes.at(0);
-  new_out_sizes.at(ndim - 1) = hw_sizes.at(1);
+  new_out_sizes.at(ndim - 2) = new_out_sizes_hw.at(0);
+  new_out_sizes.at(ndim - 1) = new_out_sizes_hw.at(1);
 
   out.virtual_resize(new_out_sizes);
   indices.virtual_resize(new_out_sizes);

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

@@ -63,7 +63,7 @@ ValueRef prepack(
     ComputeGraph& graph,
     const ValueRef vref,
     const api::GPUMemoryLayout layout) {
-  ValueRef v = graph.add_tensor(graph.get_val(vref).toTensorRef(), layout);
+  ValueRef v = graph.add_tensor_like(vref, layout);
   vTensor& t = graph.get_val(v).toTensor();
 
   api::ShaderInfo shader = get_nchw_to_image_shader(t);

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

@@ -10,9 +10,9 @@
 
 namespace vkcompute {
 
-api::utils::ivec2
-make_ivec2_int_list(ComputeGraph& graph, ValueRef vref, const bool reverse) {
-  return api::utils::make_ivec2(graph.get_val(vref).toIntList(), reverse);
+api::utils::ivec2 make_ivec2_from_list(ComputeGraph& graph, ValueRef vref) {
+  return api::utils::make_ivec2(
+      graph.get_val(vref).toIntList(), /*reverse = */ true);
 }
 
 KernelParams create_kernel_params(
@@ -22,30 +22,31 @@ KernelParams create_kernel_params(
     const ValueRef padding,
     const ValueRef dilation) {
   return {
-      make_ivec2_int_list(graph, kernel_size, /*reverse=*/true),
-      make_ivec2_int_list(graph, stride, /*reverse=*/true),
-      make_ivec2_int_list(graph, padding, /*reverse=*/true),
-      make_ivec2_int_list(graph, dilation, /*reverse=*/true),
+      make_ivec2_from_list(graph, kernel_size),
+      make_ivec2_from_list(graph, stride),
+      make_ivec2_from_list(graph, padding),
+      make_ivec2_from_list(graph, 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) {
   int64_t c = ceil_mode ? stride - 1 : 0;
   int64_t out_size =
-      (in_size + 2 * padding - dilation * (kernel - 1) - 1 + c) / stride + 1;
+      (in_size + 2 * padding - dilation * (kernel_size - 1) - 1 + c) / stride +
+      1;
   if (ceil_mode && (out_size - 1) * stride >= in_size + padding) {
     --out_size;
   }
   return out_size;
 }
 
-std::vector<int64_t> calc_hw_out_sizes(
+std::vector<int64_t> calc_out_sizes_hw(
     ComputeGraph& graph,
     const std::vector<int64_t>& in_sizes,
     const ValueRef kernel_size,
@@ -56,30 +57,28 @@ std::vector<int64_t> calc_hw_out_sizes(
   const int64_t ndim = in_sizes.size();
   std::vector<int64_t> out_sizes(2);
 
-  const auto kernel_vec =
-      make_ivec2_int_list(graph, kernel_size, /*reverse=*/false);
-  const auto stride_vec = make_ivec2_int_list(graph, stride, /*reverse=*/false);
-  const auto padding_vec =
-      make_ivec2_int_list(graph, padding, /*reverse=*/false);
-  const auto dilation_vec =
-      make_ivec2_int_list(graph, dilation, /*reverse=*/false);
+  const auto kernel_vec = make_ivec2_from_list(graph, kernel_size);
+  const auto stride_vec = make_ivec2_from_list(graph, stride);
+  const auto padding_vec = make_ivec2_from_list(graph, padding);
+  const auto dilation_vec = make_ivec2_from_list(graph, dilation);
+  const bool ceil_mode_val = graph.get_val(ceil_mode).toBool();
 
   // Height
   out_sizes.at(0) = calc_out_size(
       in_sizes.at(ndim - 2),
-      kernel_vec.data[0],
-      stride_vec.data[0],
-      padding_vec.data[0],
-      dilation_vec.data[0],
-      ceil_mode);
-  // Width
-  out_sizes.at(1) = calc_out_size(
-      in_sizes.at(ndim - 1),
       kernel_vec.data[1],
       stride_vec.data[1],
       padding_vec.data[1],
       dilation_vec.data[1],
-      ceil_mode);
+      ceil_mode_val);
+  // Width
+  out_sizes.at(1) = calc_out_size(
+      in_sizes.at(ndim - 1),
+      kernel_vec.data[0],
+      stride_vec.data[0],
+      padding_vec.data[0],
+      dilation_vec.data[0],
+      ceil_mode_val);
 
   VK_CHECK_COND(out_sizes.at(0) >= 1);
   VK_CHECK_COND(out_sizes.at(1) >= 1);

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

@@ -30,20 +30,13 @@ KernelParams create_kernel_params(
     const ValueRef padding,
     const ValueRef dilation);
 
-int64_t calc_out_size(
-    const int64_t in_size,
-    const int64_t kernel_size,
-    const int64_t stride,
-    const int64_t padding,
-    const int64_t dilation,
-    const bool ceil_mode);
-
-std::vector<int64_t> calc_hw_out_sizes(
+std::vector<int64_t> calc_out_sizes_hw(
     ComputeGraph& graph,
     const std::vector<int64_t>& in_sizes,
     const ValueRef kernel_size,
     const ValueRef stride,
     const ValueRef padding,
     const ValueRef dilation,
     const ValueRef ceil_mode);
+
 } // namespace vkcompute