diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index eed287b3a08..91da9134850 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1435,58 +1435,56 @@ def register_where(): # ============================================================================= -@update_features(exir_ops.edge.aten.index.Tensor) -def register_index_tensor(): - def _index_tensor_shapes(node: torch.fx.Node): - """(self_val, index_val) for the supported single-index form, else None.""" - self_arg = node.args[0] - indices = node.args[1] - - if not isinstance(self_arg, torch.fx.Node): - return None - self_val = self_arg.meta.get("val", None) - if self_val is None: - return None - - # Only support exactly one non-None index tensor, applied to dim 0. - if not isinstance(indices, (list, tuple)): - return None - non_none = [idx for idx in indices if idx is not None] - if len(non_none) != 1 or indices[0] is None: - return None - index_arg = non_none[0] - if not isinstance(index_arg, torch.fx.Node): - return None - index_val = index_arg.meta.get("val", None) - if index_val is None: - return None +def _index_tensor_shapes(node: torch.fx.Node): + """Return self, index, and axis for the supported form, else None.""" + self_arg = node.args[0] + indices = node.args[1] + + if not isinstance(self_arg, torch.fx.Node): + return None + self_val = self_arg.meta.get("val", None) + if self_val is None or not isinstance(indices, (list, tuple)): + return None + + non_none = [(dim, index) for dim, index in enumerate(indices) if index is not None] + if len(non_none) != 1: + return None + index_dim, index_arg = non_none[0] + if index_dim >= len(self_val.size()) or not isinstance(index_arg, torch.fx.Node): + return None + index_val = index_arg.meta.get("val", None) + if index_val is None: + return None + + return self_val, index_val, index_dim + + +def _check_index_tensor_node(node: torch.fx.Node) -> bool: + shapes = _index_tensor_shapes(node) + if shapes is None: + return False + _, index_val, _ = shapes + # The gather is expressed as "one index position per output slice", so + # the index must be 1-D. `self` may be any rank. + return len(index_val.size()) == 1 - return self_val, index_val - def check_index_tensor_node(node: torch.fx.Node) -> bool: - shapes = _index_tensor_shapes(node) - if shapes is None: - return False - _, index_val = shapes - # The gather is expressed as "one index position per output slice", so - # the index must be 1-D. `self` may be any rank: the buffer shader - # copies self's trailing dims through unchanged. - return len(index_val.size()) == 1 +def _pick_index_tensor_storage(node: torch.fx.Node): + shapes = _index_tensor_shapes(node) + # Only the buffer shader handles a higher-rank `self`. + if shapes is not None and len(shapes[0].size()) > 1: + return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER + return utils.ANY_STORAGE, utils.ANY_STORAGE - def pick_index_tensor_storage(node: torch.fx.Node): - shapes = _index_tensor_shapes(node) - # Only the buffer shader handles a higher-rank `self`; the texture - # variant still assumes the 1-D form (it reads self[idx, 0, 0, 0]). - if shapes is not None and len(shapes[0].size()) > 1: - return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER - return utils.ANY_STORAGE, utils.ANY_STORAGE +@update_features(exir_ops.edge.aten.index.Tensor) +def register_index_tensor(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, supports_resize=True, - are_node_inputs_supported_fn=check_index_tensor_node, - pick_io_storage_fn=pick_index_tensor_storage, + are_node_inputs_supported_fn=_check_index_tensor_node, + pick_io_storage_fn=_pick_index_tensor_storage, ) @@ -1500,6 +1498,7 @@ def register_arange(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, + supports_resize=True, ) diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl index 2e9377533c8..9bffc1c4132 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl @@ -23,18 +23,28 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} ${layout_declare_ubo(B, "BufferMetadata", "outp")} -${layout_declare_ubo(B, "float", "start")} -${layout_declare_ubo(B, "float", "step")} +${layout_declare_ubo(B, "uint", "start")} +${layout_declare_ubo(B, "uint", "step")} + +layout(push_constant) uniform restrict Block { + ivec2 params_are_int; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" +float decode_param(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); +} + void main() { const uint out_bufi = linear_idx_from_gid(); if (out_of_bounds(out_bufi, outp)) { return; } - t_out[out_bufi] = T(start + out_bufi * step); + const float start_val = decode_param(start, params_are_int.x); + const float step_val = decode_param(step, params_are_int.y); + t_out[out_bufi] = T(start_val + out_bufi * step_val); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl index 0a5636b300f..73c2b5e5dd6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl @@ -23,14 +23,22 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} ${layout_declare_ubo(B, "TextureMetadata", "outp")} -${layout_declare_ubo(B, "float", "start")} -${layout_declare_ubo(B, "float", "step")} +${layout_declare_ubo(B, "uint", "start")} +${layout_declare_ubo(B, "uint", "step")} + +layout(push_constant) uniform restrict Block { + ivec2 params_are_int; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")} const int packed_dim = get_packed_dim(out_layout); +float decode_param(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); +} + void main() { const ivec3 out_pos = ivec3(gl_GlobalInvocationID); @@ -44,11 +52,13 @@ void main() { // arange output is 1D, so the W dimension holds the element index. // Compute the value for each element in the texel along the packed dim. VEC4_T outtex = VEC4_T(0); + const float start_val = decode_param(start, params_are_int.x); + const float step_val = decode_param(step, params_are_int.y); int limit = min( 4, safe_idx(outp.sizes, packed_dim) - out_tidx.data[packed_dim]); for (int comp = 0; comp < limit; comp++) { int elem_idx = out_tidx.data[0]; // W index is the linear element index - outtex[comp] = VEC4_T(start + elem_idx * step).x; + outtex[comp] = VEC4_T(start_val + elem_idx * step_val).x; out_tidx.data[packed_dim]++; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl index db61e0859f2..b2497f98ccc 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl @@ -9,6 +9,7 @@ #version 450 core ${define_required_extensions("buffer", DTYPE)} +${define_required_extensions(INDEX_STORAGE, "int")} #define PRECISION ${PRECISION} @@ -22,19 +23,68 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_self", DTYPE, "buffer")} -${layout_declare_tensor(B, "r", "t_index", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_index", "int", INDEX_STORAGE)} ${layout_declare_ubo(B, "BufferMetadata", "outp")} ${layout_declare_ubo(B, "BufferMetadata", "inp")} -${layout_declare_ubo(B, "BufferMetadata", "index")} +$if INDEX_STORAGE == "buffer": + ${layout_declare_ubo(B, "BufferMetadata", "index")} +$else: + ${layout_declare_ubo(B, "TextureMetadata", "index")} + +layout(push_constant) uniform restrict Block { + ivec2 index_params; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" -// Implements aten.index.Tensor for the case where self is 1D and there is -// exactly one index tensor. Each output element is: +// Implements aten.index.Tensor with exactly one index tensor. Each output +// element is: // output[...] = self[index[...]] +${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")} +${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} +${layout_declare_spec_const(C, "int", "index_layout", "CONTIG_LAYOUT_INT")} + +int load_index(const TensorIndex out_tidx) { +$if INDEX_STORAGE == "buffer": + uint index_bufi = 0; + for (int d = 0; d < index_params.y; ++d) { + index_bufi += + stride_at(index, d) * idx_at(out_tidx, index_params.x + d); + } + return t_index[index_bufi]; +$else: + TensorIndex4D index_tidx = zero_tensor4d_idx(); + index_tidx.data.x = int(idx_at(out_tidx, index_params.x)); + if (index_params.y > 1) { + index_tidx.data.y = int(idx_at(out_tidx, index_params.x + 1)); + } + if (index_params.y > 2) { + index_tidx.data.z = int(idx_at(out_tidx, index_params.x + 2)); + } + if (index_params.y > 3) { + index_tidx.data.w = int(idx_at(out_tidx, index_params.x + 3)); + } + const TextureElementIndex index_elem = + tensor4d_idx_to_texture_element_idx_simple( + index, index_tidx, index_layout); + return texelFetch(t_index, index_elem.pos, 0)[index_elem.comp]; +} + +uint self_idx_at( + const TensorIndex out_tidx, + const int self_axis, + const uint index_value) { + if (self_axis == index_params.x) { + return index_value; + } + const int out_axis = self_axis < index_params.x + ? self_axis + : self_axis + index_params.y - 1; + return idx_at(out_tidx, out_axis); +} void main() { const uint out_bufi = linear_idx_from_gid(); @@ -45,22 +95,20 @@ void main() { // Convert output buffer index to tensor index TensorIndex out_tidx = linear_idx_to_tensor_idx(outp, out_bufi); - const uint self_rank = ndim(inp); - const uint index_rank = ndim(index); - // WHCN order places self's trailing axes before the index axes. - const uint index_axis_offset = self_rank - 1; - - uint index_bufi = 0; - for (uint d = 0; d < index_rank; ++d) { - index_bufi += - stride_at(index, d) * idx_at(out_tidx, index_axis_offset + d); - } - const int idx = t_index[index_bufi]; - - uint self_bufi = stride_at(inp, self_rank - 1) * uint(idx); - for (uint d = 0; d + 1 < self_rank; ++d) { - self_bufi += stride_at(inp, d) * idx_at(out_tidx, d); - } + const int idx = load_index(out_tidx); + + TensorIndex self_tidx; + initialize(self_tidx); + const int self_rank = int_ndim(inp); + if (self_rank > 0) self_tidx.data[0].x = self_idx_at(out_tidx, 0, uint(idx)); + if (self_rank > 1) self_tidx.data[0].y = self_idx_at(out_tidx, 1, uint(idx)); + if (self_rank > 2) self_tidx.data[0].z = self_idx_at(out_tidx, 2, uint(idx)); + if (self_rank > 3) self_tidx.data[0].w = self_idx_at(out_tidx, 3, uint(idx)); + if (self_rank > 4) self_tidx.data[1].x = self_idx_at(out_tidx, 4, uint(idx)); + if (self_rank > 5) self_tidx.data[1].y = self_idx_at(out_tidx, 5, uint(idx)); + if (self_rank > 6) self_tidx.data[1].z = self_idx_at(out_tidx, 6, uint(idx)); + if (self_rank > 7) self_tidx.data[1].w = self_idx_at(out_tidx, 7, uint(idx)); + const uint self_bufi = tensor_idx_to_linear_idx(inp, self_tidx); t_out[out_bufi] = t_self[self_bufi]; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml index ef79704203f..f4f168dfb37 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml @@ -8,7 +8,11 @@ index_tensor_buffer: parameter_names_with_default_values: DTYPE: float STORAGE: buffer + INDEX_STORAGE: buffer generate_variant_forall: + INDEX_STORAGE: + - VALUE: buffer + - VALUE: texture3d DTYPE: - VALUE: half - VALUE: float diff --git a/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl b/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl index 45aa3ed7133..5ef891e7439 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl @@ -23,16 +23,23 @@ ${define_active_storage_type(STORAGE)} layout(std430) buffer; -${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} -${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)} +${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} + +$if DYNAMIC_PARAMS: + ${layout_declare_ubo(B, "uint", "minimum")} + ${layout_declare_ubo(B, "uint", "maximum")} layout(push_constant) uniform restrict Block { $if STORAGE == "buffer": int numel; $else: ivec4 out_limits; -float minimum; -float maximum; +$if DYNAMIC_PARAMS: + ivec2 bounds_are_int; +$else: + float minimum; + float maximum; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -40,6 +47,11 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" #include "activations.h" +$if DYNAMIC_PARAMS: + float decode_bound(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); + } + #ifdef USING_BUFFER void main() { @@ -48,7 +60,13 @@ void main() { return; } - float in_val = float(t_in[i]); +$if DYNAMIC_PARAMS: + const T in_val = T(t_in[i]); + const T minimum_val = T(decode_bound(minimum, bounds_are_int.x)); + const T maximum_val = T(decode_bound(maximum, bounds_are_int.y)); + t_out[i] = T(op(in_val, minimum_val, maximum_val)); +$else: + const float in_val = float(t_in[i]); t_out[i] = T(op(in_val, minimum, maximum)); } @@ -62,6 +80,11 @@ void main() { } VEC4_T in_texel = texelFetch(t_in, pos, 0); +$if DYNAMIC_PARAMS: + const VEC4_T minimum_val = VEC4_T(decode_bound(minimum, bounds_are_int.x)); + const VEC4_T maximum_val = VEC4_T(decode_bound(maximum, bounds_are_int.y)); + imageStore(t_out, pos, op(in_texel, minimum_val, maximum_val)); +$else: imageStore(t_out, pos, VEC4_T(op(in_texel, minimum, maximum))); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml index 46d12806149..0331a15fde6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml @@ -3,6 +3,7 @@ unary_op: OPERATOR: clamp(X, A, B) DTYPE: float STORAGE: texture3d + DYNAMIC_PARAMS: false generate_variant_forall: DTYPE: - VALUE: half @@ -18,6 +19,13 @@ unary_op: - NAME: clamp_int32 OPERATOR: clamp(X, A, B) DTYPE: int32 + - NAME: clamp_dynamic_int32 + OPERATOR: clamp(X, A, B) + DTYPE: int32 + DYNAMIC_PARAMS: true + - NAME: clamp_dynamic + OPERATOR: clamp(X, A, B) + DYNAMIC_PARAMS: true - NAME: cos OPERATOR: cos(X) - NAME: exp diff --git a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp index f635c9282f2..839b94f5e75 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp @@ -15,6 +15,8 @@ #include +#include + namespace vkcompute { void resize_arange_node( @@ -23,18 +25,22 @@ void resize_arange_node( const std::vector& extra_args) { const ValueRef out = args.at(0).refs.at(0); - int start_val = 0; - int step_val = 1; + double start_val = 0.0; + double step_val = 1.0; if (!graph->val_is_none(extra_args.at(0))) { - start_val = graph->extract_scalar(extra_args.at(0)); + start_val = graph->extract_scalar(extra_args.at(0)); } - const int end_val = graph->extract_scalar(extra_args.at(1)); + const double end_val = graph->extract_scalar(extra_args.at(1)); if (!graph->val_is_none(extra_args.at(2))) { - step_val = graph->extract_scalar(extra_args.at(2)); + step_val = graph->extract_scalar(extra_args.at(2)); } + VK_CHECK_COND(step_val != 0.0, "arange: step must be nonzero"); + const double range_size = (end_val - start_val) / step_val; + VK_CHECK_COND( + range_size >= 0.0, "arange: bounds are inconsistent with step sign"); const std::vector out_sizes = { - utils::div_up(end_val - start_val, step_val)}; + static_cast(std::ceil(range_size))}; graph->virtual_resize(out, out_sizes); } @@ -55,39 +61,35 @@ void check_arange_input( } } +vkapi::BufferBindInfo get_arange_param_buffer( + ComputeGraph& graph, + const ValueRef value, + const float default_value) { + if (graph.val_is_symint(value)) { + return graph.get_or_create_int_param_buffer(value); + } + return graph.create_params_buffer( + graph.extract_scalar_or(value, default_value)); +} + void add_arange_node( ComputeGraph& graph, const ValueRef start, const ValueRef end, const ValueRef step, const ValueRef out) { - float start_val = 0.0f; - float step_val = 1.0f; - if (graph.val_is_none(end)) { VK_THROW("arange: end must be specified!"); } - if (!graph.val_is_none(start)) { - if (graph.val_is_int(start)) { - start_val = static_cast(graph.extract_scalar(start)); - } else { - start_val = graph.extract_scalar(start); - } - } - if (!graph.val_is_none(step)) { - if (graph.val_is_int(step)) { - step_val = static_cast(graph.extract_scalar(step)); - } else { - step_val = graph.extract_scalar(step); - } - } - std::string kernel_name("arange"); kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); + const utils::ivec2 params_are_int = { + graph.val_is_symint(start) ? 1 : 0, graph.val_is_symint(step) ? 1 : 0}; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), @@ -97,10 +99,10 @@ void add_arange_node( {{out, vkapi::kWrite}}, // Shader params buffers {graph.meta_ubo(out), - graph.create_params_buffer(start_val), - graph.create_params_buffer(step_val)}, + get_arange_param_buffer(graph, start, 0.0f), + get_arange_param_buffer(graph, step, 1.0f)}, // Push Constants - {}, + {PushConstantDataInfo(¶ms_are_int, sizeof(params_are_int))}, // Specialization Constants {graph.hashed_layout_of(out)}, // Resize Args diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp index ddd8e8994b1..f490f60f75c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp @@ -9,6 +9,7 @@ #include #include +#include #include @@ -18,19 +19,31 @@ void resize_index_tensor_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { - (void)resize_args; const ValueRef out = args.at(0).refs.at(0); const ValueRef self = args.at(1).refs.at(0); const ValueRef index = args.at(1).refs.at(1); - // aten.index.Tensor with a single index tensor gathers along dim 0, so - // out.sizes = index.sizes ++ self.sizes[1:] - // Using the index's sizes alone is only correct when self is 1-D; for any - // higher-rank self it also changes the tensor's RANK, which virtual_resize - // rejects outright ("new sizes cannot modify the dimensionality"). + int64_t index_dim = -1; + { + const ValueListPtr indices = graph->get_value_list(resize_args.at(0)); + for (size_t dim = 0; dim < indices->size(); ++dim) { + if (!graph->val_is_none(indices->at(dim))) { + index_dim = utils::safe_downcast(dim); + break; + } + } + } + VK_CHECK_COND(index_dim >= 0, "index.Tensor: an index tensor is required"); + const std::vector self_sizes = graph->sizes_of(self); - std::vector out_sizes = graph->sizes_of(index); - out_sizes.insert(out_sizes.end(), self_sizes.begin() + 1, self_sizes.end()); + const std::vector index_sizes = graph->sizes_of(index); + std::vector out_sizes; + out_sizes.reserve(self_sizes.size() + index_sizes.size() - 1); + out_sizes.insert( + out_sizes.end(), self_sizes.begin(), self_sizes.begin() + index_dim); + out_sizes.insert(out_sizes.end(), index_sizes.begin(), index_sizes.end()); + out_sizes.insert( + out_sizes.end(), self_sizes.begin() + index_dim + 1, self_sizes.end()); graph->virtual_resize(out, out_sizes); } @@ -39,14 +52,24 @@ void add_index_tensor_node( ComputeGraph& graph, const ValueRef self, const ValueRef index, + const int64_t index_dim, + const ValueRef indices_list_ref, const ValueRef out) { std::string kernel_name = "index_tensor"; kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + if (graph.is_buffer_storage(out)) { + add_storage_type_suffix(kernel_name, graph.storage_type_of(index)); + } add_dtype_suffix(kernel_name, graph.dtype_of(out)); vkapi::ParamsBindList param_ubos = { graph.meta_ubo(out), graph.meta_ubo(self), graph.meta_ubo(index)}; + const utils::ivec2 index_params = { + utils::safe_downcast(graph.dim_of(self) - 1 - index_dim), + utils::safe_downcast(graph.dim_of(index))}; + std::vector push_constants = { + PushConstantDataInfo(&index_params, sizeof(index_params))}; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -58,11 +81,13 @@ void add_index_tensor_node( // Shader params buffers param_ubos, // Push Constants - {}, + push_constants, // Specialization Constants - {graph.hashed_layout_of(out), graph.hashed_layout_of(self)}, + {graph.hashed_layout_of(out), + graph.hashed_layout_of(self), + graph.hashed_layout_of(index)}, // Resize Args - {}, + {indices_list_ref}, // Resizing Logic resize_index_tensor_node)); } @@ -72,14 +97,27 @@ void index_tensor(ComputeGraph& graph, const std::vector& args) { ValueRef indices_list_ref = args[1]; ValueRef out = args[2]; - ValueListPtr indices_list = graph.get_value_list(indices_list_ref); + ValueRef index = -1; + int64_t index_dim = -1; + { + const ValueListPtr indices_list = graph.get_value_list(indices_list_ref); + for (size_t dim = 0; dim < indices_list->size(); ++dim) { + const ValueRef candidate = indices_list->at(dim); + if (graph.val_is_none(candidate)) { + continue; + } + VK_CHECK_COND( + index_dim < 0, "index.Tensor: only one index tensor is supported"); + index = candidate; + index_dim = utils::safe_downcast(dim); + } + } + VK_CHECK_COND(index_dim >= 0, "index.Tensor: an index tensor is required"); VK_CHECK_COND( - indices_list->size() == 1, - "index.Tensor: only one index tensor is supported"); - - ValueRef index = indices_list->at(0); + index_dim < graph.dim_of(self), + "index.Tensor: index dimension is invalid"); - add_index_tensor_node(graph, self, index, out); + add_index_tensor_node(graph, self, index, index_dim, indices_list_ref, out); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp index 6a50cb2f6a9..d17f57774f7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp @@ -69,6 +69,46 @@ void add_unary_op_node( resize_unary_op_node)); } +void add_dynamic_clamp_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef min, + const ValueRef max, + const ValueRef out) { + std::string kernel_name("clamp_dynamic"); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + + const bool output_is_int = graph.dtype_of(out) == vkapi::kInt; + const utils::ivec2 bounds_are_int = { + output_is_int || graph.val_is_symint(min) ? 1 : 0, + output_is_int || graph.val_is_symint(max) ? 1 : 0}; + const vkapi::BufferBindInfo min_param = bounds_are_int[0] + ? graph.get_or_create_int_param_buffer( + min, std::numeric_limits::min()) + : graph.create_params_buffer(graph.extract_scalar_or( + min, -std::numeric_limits::infinity())); + const vkapi::BufferBindInfo max_param = bounds_are_int[1] + ? graph.get_or_create_int_param_buffer( + max, std::numeric_limits::max()) + : graph.create_params_buffer(graph.extract_scalar_or( + max, std::numeric_limits::infinity())); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_gwg, + default_pick_lwg, + {{out, vkapi::kWrite}, {in, vkapi::kRead}}, + {min_param, max_param}, + {graph.is_buffer_storage(out) ? graph.numel_pc_of(out) + : graph.logical_limits_pc_of(out), + PushConstantDataInfo(&bounds_are_int, sizeof(bounds_are_int))}, + {}, + {}, + resize_unary_op_node)); +} + float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) { if (!graph.val_is_none(val)) { return graph.extract_scalar(val); @@ -85,6 +125,10 @@ float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) { #define DEFINE_CLAMP_FN(op_name) \ void op_name(ComputeGraph& graph, const std::vector& args) { \ + if (graph.val_is_symint(args[1]) || graph.val_is_symint(args[2])) { \ + return add_dynamic_clamp_node( \ + graph, args[0], args[1], args[2], args[3]); \ + } \ return add_unary_op_node( \ graph, \ args[0], \ diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index c6915d37684..34a4f12f62c 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -13,6 +13,7 @@ import executorch.backends.vulkan.test.utils as test_utils import torch +import torch.nn.functional as F from executorch.backends.transforms.convert_dtype_pass import I64toI32 from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.vulkan.vulkan_preprocess import VulkanBackend @@ -538,6 +539,20 @@ def forward(self, x): self.lower_module_and_test_output(ClampModule(), sample_inputs) + def test_vulkan_backend_dynamic_float_clamp(self): + class ClampModule(torch.nn.Module): + def forward(self, x): + return torch.clamp(x, max=x.shape[0]) + + sample_inputs = (torch.arange(32).reshape(8, 4).float(),) + length = Dim("length", min=2, max=16) + self.lower_module_and_test_output( + ClampModule(), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.arange(12).reshape(3, 4).float(),)], + ) + def test_vulkan_backend_cos(self): class CosModule(torch.nn.Module): def __init__(self): @@ -1625,6 +1640,84 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_index_tensor_nonzero_axis(self): + class IndexTensorModule(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + self.index = torch.tensor([0, 2]) + + def forward(self, x): + indices = [slice(None)] * x.dim() + indices[self.dim] = self.index + return x[tuple(indices)] + + sample_inputs = (torch.arange(24).reshape(1, 3, 8).float(),) + for dim in (1, 2): + self.lower_module_and_test_output( + IndexTensorModule(dim), + sample_inputs, + ) + + def test_vulkan_backend_dynamic_replicate_pad_time_reduction(self): + class TimeReductionModule(torch.nn.Module): + def forward(self, x): + padded_frames = 8 * ((x.shape[1] + 7) // 8) + x = F.pad( + x, + (0, 0, 0, padded_frames - x.shape[1]), + mode="replicate", + ) + return x.view(x.shape[0], -1, 640) + + sample_inputs = (torch.randn(1, 24, 80),) + frames = Dim("frames", min=1, max=24) + self.lower_module_and_test_output( + TimeReductionModule(), + sample_inputs, + dynamic_shapes={"x": {1: frames}}, + test_inputs=[ + (torch.randn(1, 8, 80),), + (torch.randn(1, 9, 80),), + (torch.randn(1, 17, 80),), + ], + ) + + def test_vulkan_backend_dynamic_arange_float_step(self): + class ArangeModule(torch.nn.Module): + def __init__(self, end_scale, step): + super().__init__() + self.end_scale = end_scale + self.step = step + + def forward(self, x): + return torch.arange(0, self.end_scale * x.shape[0], self.step) + + sample_inputs = (torch.randn(8),) + length = Dim("length", min=2, max=16) + for end_scale, step in ((1, 0.5), (-1, -0.5)): + with self.subTest(end_scale=end_scale, step=step): + self.lower_module_and_test_output( + ArangeModule(end_scale, step), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.randn(3),), (torch.randn(7),)], + ) + + def test_vulkan_backend_dynamic_arange_start(self): + class ArangeModule(torch.nn.Module): + def forward(self, x): + return torch.arange(x.shape[0], 32, 2) + + sample_inputs = (torch.randn(8),) + length = Dim("length", min=2, max=16) + self.lower_module_and_test_output( + ArangeModule(), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.randn(3),), (torch.randn(15),)], + ) + def test_vulkan_backend_arange_int(self): class ArangeModule(torch.nn.Module): def __init__(self, input):