From 8dc3bcb032e1fdcb362403b76de6b62a1e0da5b2 Mon Sep 17 00:00:00 2001 From: Alex Dean Date: Wed, 6 Aug 2025 17:58:57 -0700 Subject: [PATCH] [ET-VK] Fix Build Errors in Vulkan Backend --- backends/vulkan/op_registry.py | 28 ++++++------------- .../vulkan/runtime/graph/ops/impl/Reduce.cpp | 8 +++--- 2 files changed, 12 insertions(+), 24 deletions(-) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index b3dd86e1387..22a93ec0e2b 100644 --- a/backends/vulkan/op_registry.py +++ b/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 diff --git a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp index d4f0b1e29c8..6ad1d7f371d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp @@ -37,8 +37,8 @@ void resize_reduce2d_node( ComputeGraph* graph, const std::vector& args, const std::vector& 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 dims_list = @@ -46,10 +46,10 @@ void resize_reduce2d_node( int32_t reduce_dim1_nchw = dims_list[0]; int32_t reduce_dim2_nchw = dims_list[1]; - std::vector new_sizes = in->sizes(); + std::vector 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(