[ET-VK][ez] Introduce a graph config setting to force resize functions to execute

Pull Request resolved: #15158

Title says it all! A few months back, a mechanism was introduced where an `ExecuteNode` would not call an operator's resize function if none of the arguments were updated.

However, this creates a blind spot during testing where the resize function of operators are not tested since the generated operator tests do not modify input sizes.

To address this, add a way to force the resize function to be called during testing.
ghstack-source-id: 317683546
@exported-using-ghexport

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

Author: ssjia

backends/vulkan/runtime/graph/GraphConfig.cpp

@@ -65,6 +65,7 @@ GraphConfig::GraphConfig() {
   local_wg_size_override = {};
 
   expect_dynamic_shapes = false;
+  force_resize = false;
 
   external_adapter = nullptr;
 }

backends/vulkan/runtime/graph/GraphConfig.h

@@ -35,6 +35,9 @@ struct GraphConfig final {
 
   // Whether or not the ComputeGraph should expect input shapes to be dynamic
   bool expect_dynamic_shapes;
+  // Used for testing/debugging only. Forces ExecuteNode to trigger the resize
+  // function even if none of the inputs have been updated.
+  bool force_resize = false;
 
   // Execution properties that determine specifics re: how command buffer
   // submission is handled, etc. 0 means this field is not set.

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

@@ -21,9 +21,10 @@ ExecuteNode::ExecuteNode(
       name_(name) {}
 
 bool ExecuteNode::trigger_resize(ComputeGraph* graph) {
-  const bool any_arg_updated = was_any_arg_updated(graph);
-  if (resize_fn_ && any_arg_updated) {
+  bool any_arg_updated = was_any_arg_updated(graph);
+  if (resize_fn_ && (any_arg_updated || graph->graphconfig().force_resize)) {
     resize_fn_(graph, args_, resize_args_);
+    any_arg_updated = true;
   }
   return any_arg_updated;
 }

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

@@ -59,7 +59,11 @@ void resize_pool2d_node(
 
   if (is_max_pool2d) {
     const ValueRef indices = args.at(0).refs.at(1);
-    graph->virtual_resize(indices, new_out_sizes);
+    // For max_pool2d variant, indices tensor will be a 0-dim tensor - only
+    // resize the indices tensor if this is not the case.
+    if (graph->sizes_of(indices).size() > 0) {
+      graph->virtual_resize(indices, new_out_sizes);
+    }
   }
 }
 

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

@@ -54,12 +54,21 @@ void resize_unsqueeze_node(
   const ValueRef in = args.at(1).refs.at(0);
   const ValueRef dims_ref = extra_args.at(0);
 
-  const IntListPtr dims = graph->get_int_list(dims_ref);
+  std::vector<int64_t> dims_vec;
+  if (graph->is_scalar_or_none(dims_ref)) {
+    // Handle scalar case
+    int64_t dim = graph->extract_scalar<int64_t>(dims_ref);
+    dims_vec.push_back(dim);
+  } else {
+    // Handle list case
+    const IntListPtr dims = graph->get_int_list(dims_ref);
+    dims_vec.assign(dims->begin(), dims->end());
+  }
 
   std::vector<int64_t> out_sizes = graph->sizes_of(in);
 
   // Insert singleton dimensions at the specified positions
-  for (auto dim : *dims) {
+  for (auto dim : dims_vec) {
     int64_t d = dim;
     if (d < 0) {
       d += static_cast<int64_t>(out_sizes.size()) + 1;

backends/vulkan/test/op_tests/utils/gen_correctness_vk.py

@@ -34,6 +34,7 @@ class GeneratedOpsTest_{op_name} : public ::testing::TestWithParam< ::std::tuple
     std::tie(test_dtype, default_storage_type, default_memory_layout) = GetParam();
     config.set_storage_type_override(default_storage_type);
     config.set_memory_layout_override(default_memory_layout);
+    config.force_resize = true;
     graph = new ComputeGraph(config);
 
     if (test_dtype == at::kHalf) {{