diff --git a/backends/vulkan/runtime/graph/ops/glsl/upsample.glsl b/backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.glsl similarity index 77% rename from backends/vulkan/runtime/graph/ops/glsl/upsample.glsl rename to backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.glsl index 7664c2a9ee..8ab455a55a 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/upsample.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.glsl @@ -20,6 +20,8 @@ layout(std430) buffer; ${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} ${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)} ${layout_declare_ubo(2, "ivec3", "out_limits")} +${layout_declare_ubo(3, "ivec2", "input_size")} +${layout_declare_ubo(4, "vec2", "rev_scales")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -30,6 +32,8 @@ void main() { return; } - VEC4_T in_texel = texelFetch(t_in, pos, 0); + const ivec2 ipos = clamp(ivec2(pos.xy * rev_scales), ivec2(0), input_size); + + VEC4_T in_texel = texelFetch(t_in, ivec3(ipos, pos.z), 0); imageStore(t_out, pos, in_texel); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/upsample.yaml b/backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.yaml similarity index 89% rename from backends/vulkan/runtime/graph/ops/glsl/upsample.yaml rename to backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.yaml index af0dbb6152..89b873c02c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/upsample.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/upsample_nearest2d.yaml @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -upsample: +upsample_nearest2d: parameter_names_with_default_values: NDIM: 3 DTYPE: float @@ -15,4 +15,4 @@ upsample: - VALUE: half - VALUE: float shader_variants: - - NAME: upsample + - NAME: upsample_nearest2d diff --git a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp index 2ba80b1284..e8fc7bbeda 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp @@ -10,33 +10,56 @@ #include -#include #include #include namespace vkcompute { -void resize_upsample_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& extra_args) { - (void)graph; - (void)args; - (void)extra_args; -} +// Executorch-Vulkan framework to add node +// Args: +// in: will be converted from NCHW input tensor to 3D ARGB representation in +// openGL (via Executorch) output_sizes: optional 2D array of targetting +// output size of H and W dimensions. >= input sizes; -void add_upsample_node( +// will be computed if only given the scale_factors. +// scale_factors: optional 2D array of scale factors for H and W dimensions. +// Will be computed if only given the output_sizes. +void add_upsample_nearest2d_node( ComputeGraph& graph, const ValueRef in, + const ValueRef output_sizes, + const ValueRef scale_factors, const ValueRef out) { - ValueRef arg = prepack_if_tensor_ref(graph, in); + // TODO(T190297757) add supports for output_sizes + if (graph.val_is_none(output_sizes) && graph.val_is_none(scale_factors)) { + VK_THROW( + "Invalid input, must provide either output_sizes or scale_factors"); + } + if (!graph.val_is_none(output_sizes) && !graph.val_is_none(scale_factors)) { + VK_THROW( + "Invalid input, must provide ONLY one of output_sizes or scale_factors"); + } + auto scales = graph.get_double_list(scale_factors); + + ValueRef arg_in = prepack_if_tensor_ref(graph, in); + + vTensorPtr t_in = graph.get_tensor(in); + api::utils::uvec3 input_sizes = t_in->image_extents(); + + api::utils::ivec2 input_size = { + api::utils::safe_downcast(input_sizes.data[0]), + api::utils::safe_downcast(input_sizes.data[1])}; + // Reverse scale factors that pre-computed before GLSL. + api::utils::vec2 rev_scales = { + api::utils::safe_downcast(1.0 / scales->at(1)), + api::utils::safe_downcast(1.0 / scales->at(0))}; vTensorPtr t_out = graph.get_tensor(out); api::utils::uvec3 global_size = t_out->image_extents(); api::utils::uvec3 local_size = adaptive_work_group_size(global_size); - std::string kernel_name("upsample"); + std::string kernel_name("upsample_nearest2d"); kernel_name.reserve(kShaderNameReserve); add_dtype_suffix(kernel_name, *t_out); @@ -47,17 +70,18 @@ void add_upsample_node( global_size, local_size, // Inputs and Outputs - {{out, api::MemoryAccessType::WRITE}, {arg, api::MemoryAccessType::READ}}, + {{out, api::MemoryAccessType::WRITE}, + {arg_in, api::MemoryAccessType::READ}}, // Shader params buffers - {t_out->texture_limits_ubo()}, + {t_out->texture_limits_ubo(), + graph.create_params_buffer(input_size), + graph.create_params_buffer(rev_scales)}, // Specialization Constants - {}, - // Resizing Logic - resize_upsample_node)); + {})); } void upsample(ComputeGraph& graph, const std::vector& args) { - return add_upsample_node(graph, args[0], args[3]); + return add_upsample_nearest2d_node(graph, args[0], args[1], args[2], args[3]); } REGISTER_OPERATORS { diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py index 89443c98bc..ff74bca578 100644 --- a/backends/vulkan/test/op_tests/cases.py +++ b/backends/vulkan/test/op_tests/cases.py @@ -244,8 +244,12 @@ def get_native_layer_norm_inputs(): def get_upsample_inputs(): test_suite = VkTestSuite( [ - # TODO(dixu): implement the basic upsample logic to have a meaningful test + # (input tensor shape, output 2D image size (H, W), output scaling factors) ((2, 2, 2, 2), None, [1, 1]), + ((1, 1, 2, 2), None, [2, 2]), + ((1, 1, 2, 2), None, [2, 4]), + ((1, 1, 2, 2), None, [4, 2]), + # TODO(T190297757) add supports for output_sizes ] ) return test_suite