1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
10
11 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
12
13 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
14
15 #include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
16
17 namespace vkcompute {
18
resize_upsample_nearest2d_node(ComputeGraph * graph,const std::vector<ArgGroup> & args,const std::vector<ValueRef> & extra_args)19 void resize_upsample_nearest2d_node(
20 ComputeGraph* graph,
21 const std::vector<ArgGroup>& args,
22 const std::vector<ValueRef>& extra_args) {
23 vTensorPtr out = graph->get_tensor(args[0].refs[0]);
24 vTensorPtr self = graph->get_tensor(args[1].refs[0]);
25 std::vector<int64_t> out_sizes = self->sizes(); // NCHW
26
27 const ValueRef output_sizes = extra_args[0]; // HW
28 const ValueRef scale_factors = extra_args[1]; // HW
29 if (!graph->val_is_none(output_sizes)) {
30 IntListPtr output_size_ref = graph->get_int_list(output_sizes);
31 out_sizes.at(2) = output_size_ref->at(0);
32 out_sizes.at(3) = output_size_ref->at(1);
33 } else {
34 DoubleListPtr scales = graph->get_double_list(scale_factors);
35 out_sizes.at(2) *= scales->at(0);
36 out_sizes.at(3) *= scales->at(1);
37 }
38
39 out->virtual_resize(out_sizes);
40 }
41
42 // ExecuTorch-Vulkan framework to add node
43 // Args:
44 // in: will be converted from NCHW input tensor to 3D ARGB representation in
45 // openGL (via ExecuTorch) output_sizes: optional 2D array of targetting
46 // output size of H and W dimensions. >= input sizes;
47
48 // will be computed if only given the scale_factors.
49 // scale_factors: optional 2D array of scale factors for H and W dimensions.
50 // Will be computed if only given the output_sizes.
add_upsample_nearest2d_node(ComputeGraph & graph,const ValueRef in,const ValueRef output_sizes,const ValueRef scale_factors,const ValueRef out)51 void add_upsample_nearest2d_node(
52 ComputeGraph& graph,
53 const ValueRef in,
54 const ValueRef output_sizes,
55 const ValueRef scale_factors,
56 const ValueRef out) {
57 if (graph.val_is_none(output_sizes) && graph.val_is_none(scale_factors)) {
58 VK_THROW(
59 "Invalid input, must provide either output_sizes or scale_factors");
60 }
61 if (!graph.val_is_none(output_sizes) && !graph.val_is_none(scale_factors)) {
62 VK_THROW(
63 "Invalid input, must provide ONLY one of output_sizes or scale_factors");
64 }
65
66 vTensorPtr t_in = graph.get_tensor(in);
67 utils::uvec3 input_sizes = t_in->logical_limits();
68
69 utils::ivec2 input_size = {
70 utils::safe_downcast<int32_t>(input_sizes[0]),
71 utils::safe_downcast<int32_t>(input_sizes[1])};
72 utils::vec2 rev_scales = {
73 utils::safe_downcast<float>(1.0), utils::safe_downcast<float>(1.0)};
74
75 // Reverse scale factors that pre-computed before GLSL.
76 if (!graph.val_is_none(output_sizes)) {
77 auto output_size_ref = graph.get_int_list(output_sizes);
78 rev_scales = {
79 utils::safe_downcast<float>(
80 (float)input_size[0] / output_size_ref->at(1)),
81 utils::safe_downcast<float>(
82 (float)input_size[1] / output_size_ref->at(0))};
83
84 } else {
85 auto scales = graph.get_double_list(scale_factors);
86 rev_scales = {
87 utils::safe_downcast<float>(1.0 / scales->at(1)),
88 utils::safe_downcast<float>(1.0 / scales->at(0))};
89 }
90
91 vTensorPtr t_out = graph.get_tensor(out);
92
93 std::string kernel_name("upsample_nearest2d");
94 kernel_name.reserve(kShaderNameReserve);
95 add_dtype_suffix(kernel_name, *t_out);
96
97 graph.execute_nodes().emplace_back(new DispatchNode(
98 graph,
99 VK_KERNEL_FROM_STR(kernel_name),
100 graph.create_global_wg_size(out),
101 graph.create_local_wg_size(out),
102 // Inputs and Outputs
103 {{out, vkapi::MemoryAccessType::WRITE},
104 {in, vkapi::MemoryAccessType::READ}},
105 // Shader params buffers
106 {t_out->logical_limits_ubo(),
107 graph.create_params_buffer(input_size),
108 graph.create_params_buffer(rev_scales)},
109 // Specialization Constants
110 {},
111 resize_upsample_nearest2d_node,
112 {output_sizes, scale_factors}));
113 }
114
upsample(ComputeGraph & graph,const std::vector<ValueRef> & args)115 void upsample(ComputeGraph& graph, const std::vector<ValueRef>& args) {
116 return add_upsample_nearest2d_node(graph, args[0], args[1], args[2], args[3]);
117 }
118
119 REGISTER_OPERATORS {
120 VK_REGISTER_OP(aten.upsample_nearest2d.vec, upsample);
121 }
122
123 } // namespace vkcompute
124