[ET-VK] Fix Build Errors in Vulkan Backend

Author: alexdean08

backends/vulkan/op_registry.py

@@ -384,26 +384,14 @@ def check_reduce_node(node: torch.fx.Node) -> bool:
                 memory_layout = utils.get_node_memory_layout(node)
 
                 # If we have memory layout information, check if any dimension in dim_list corresponds to a packed dimension
-                if memory_layout is not None:
-                    for dim in dim_list:
-                        # For WIDTH_PACKED layout, dimension 3 (W) is packed
-                        # For HEIGHT_PACKED layout, dimension 2 (H) is packed
-                        # For CHANNELS_PACKED layout, dimension 1 (C) is packed
-                        if (
-                            (
-                                memory_layout == VkMemoryLayout.TENSOR_WIDTH_PACKED
-                                and dim == 3
-                            )
-                            or (
-                                memory_layout == VkMemoryLayout.TENSOR_HEIGHT_PACKED
-                                and dim == 2
-                            )
-                            or (
-                                memory_layout == VkMemoryLayout.TENSOR_CHANNELS_PACKED
-                                and dim == 1
-                            )
-                        ):
-                            return False
+                if (
+                    memory_layout is not None
+                    and memory_layout != VkMemoryLayout.DEFAULT_LAYOUT
+                ):
+                    # For now only default layout is supported for 2D reduction.
+                    # Because we can't determine if the input is NCHW or NHWC here,
+                    # assume the reduction dimension is packed so we cannot support it.
+                    return False
             except (AssertionError, KeyError, AttributeError):
                 # If we can't get memory layout information, we'll assume the dims aren't packed
                 pass

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

@@ -37,19 +37,19 @@ void resize_reduce2d_node(
     ComputeGraph* graph,
     const std::vector<ArgGroup>& args,
     const std::vector<ValueRef>& resize_args) {
-  vTensorPtr out = graph->get_tensor(args[0].refs[0]);
-  vTensorPtr in = graph->get_tensor(args[1].refs[0]);
+  const ValueRef out = args.at(0).refs.at(0);
+  const ValueRef in = args.at(1).refs.at(0);
 
   // Extract the dimensions to reduce over
   const std::vector<int64_t> dims_list =
       graph->extract_int_or_symint_list(resize_args.at(0));
   int32_t reduce_dim1_nchw = dims_list[0];
   int32_t reduce_dim2_nchw = dims_list[1];
 
-  std::vector<int64_t> new_sizes = in->sizes();
+  std::vector<int64_t> new_sizes = graph->sizes_of(in);
   new_sizes.at(normalize(reduce_dim1_nchw, new_sizes.size())) = 1;
   new_sizes.at(normalize(reduce_dim2_nchw, new_sizes.size())) = 1;
-  out->virtual_resize(new_sizes);
+  graph->virtual_resize(out, new_sizes);
 }
 
 utils::uvec3 reduce_global_wg_size(