From c73dbe0791b40e946d3ee96367d043da4a561c2b Mon Sep 17 00:00:00 2001 From: ssjia Date: Tue, 10 Mar 2026 10:00:49 -0700 Subject: [PATCH] [ET-VK] Add ANY_STORAGE support to arange and full Split arange and full shaders into separate buffer and texture variants using BufferMetadata/TextureMetadata with indexing.glslh for layout-agnostic access. Update the C++ dispatch code to use add_storage_type_suffix and meta_ubo instead of sizes_ubo so both storage types are dispatched correctly. Update Full.cpp to use extract_int_or_symint_list() for symint support with dynamic shapes. Update op_registry.py to advertise ANY_STORAGE for both arange and full. Add kBuffer storage type coverage to the test suites. Differential Revision: [D95970169](https://our.internmc.facebook.com/intern/diff/D95970169/) [ghstack-poisoned] --- backends/vulkan/op_registry.py | 4 +- .../vulkan/runtime/graph/ops/glsl/arange.glsl | 39 -------------- .../runtime/graph/ops/glsl/arange_buffer.glsl | 38 ++++++++++++++ .../glsl/{arange.yaml => arange_buffer.yaml} | 10 ++-- .../graph/ops/glsl/arange_texture.glsl | 52 +++++++++++++++++++ .../graph/ops/glsl/arange_texture.yaml | 16 ++++++ .../runtime/graph/ops/glsl/full_buffer.glsl | 36 +++++++++++++ .../runtime/graph/ops/glsl/full_buffer.yaml | 18 +++++++ .../ops/glsl/{full.glsl => full_texture.glsl} | 27 +++++----- .../ops/glsl/{full.yaml => full_texture.yaml} | 6 +-- .../vulkan/runtime/graph/ops/impl/Arange.cpp | 3 +- .../vulkan/runtime/graph/ops/impl/Full.cpp | 7 +-- backends/vulkan/test/op_tests/cases.py | 7 +++ 13 files changed, 195 insertions(+), 68 deletions(-) delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/arange.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl rename backends/vulkan/runtime/graph/ops/glsl/{arange.yaml => arange_buffer.yaml} (78%) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/arange_texture.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/full_buffer.yaml rename backends/vulkan/runtime/graph/ops/glsl/{full.glsl => full_texture.glsl} (59%) rename backends/vulkan/runtime/graph/ops/glsl/{full.yaml => full_texture.yaml} (88%) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index d68f62fa0e7..efe2bc98070 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1250,7 +1250,7 @@ def check_index_tensor_node(node: torch.fx.Node) -> bool: @update_features(exir_ops.edge.aten.arange.start_step) def register_arange(): return OpFeatures( - inputs_storage=utils.CHANNELS_PACKED_TEXTURE, + inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, ) @@ -1286,7 +1286,7 @@ def register_constant_pad_nd(): ) def register_full_cpp_ops(): return OpFeatures( - inputs_storage=utils.CHANNELS_PACKED_TEXTURE, + inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_BOOL_T, ) diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange.glsl deleted file mode 100644 index 8b1841888ad..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/arange.glsl +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#version 450 core - -#define PRECISION ${PRECISION} - -#define VEC4_T ${texel_type(DTYPE)} - -layout(std430) buffer; - -#include "indexing_utils.h" - -${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} -${layout_declare_ubo(1, "ivec4", "sizes")} -${layout_declare_ubo(2, "float", "start")} -${layout_declare_ubo(3, "float", "step")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -layout(constant_id = 3) const int packed_dim = C_DIM; - -void main() { - const ivec3 pos = ivec3(gl_GlobalInvocationID); - const ivec4 idx = to_tensor_idx(pos, sizes, packed_dim); - - if (pos_out_of_bounds(pos, sizes, packed_dim)) { - return; - } - - VEC4_T outtex = VEC4_T(start + pos.x * step, 0, 0, 0); - - imageStore(t_out, pos, outtex); -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl new file mode 100644 index 00000000000..906ed91a466 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl @@ -0,0 +1,38 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#version 450 core + +${define_required_extensions("buffer", DTYPE)} + +#define PRECISION ${PRECISION} + +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} + +layout(std430) buffer; + +#include "indexing.glslh" + +${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(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const uint out_bufi = gl_GlobalInvocationID.x; + if (out_of_bounds(out_bufi, outp)) { + return; + } + + t_out[out_bufi] = T(start + out_bufi * step); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange.yaml b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.yaml similarity index 78% rename from backends/vulkan/runtime/graph/ops/glsl/arange.yaml rename to backends/vulkan/runtime/graph/ops/glsl/arange_buffer.yaml index 37b2027db85..d53d0387788 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/arange.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.yaml @@ -4,16 +4,14 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -arange: +arange_buffer: parameter_names_with_default_values: - NDIM: 3 - DTYPE: int32 - STORAGE: texture3d - PACKING: C_packed + DTYPE: float + STORAGE: buffer generate_variant_forall: DTYPE: - VALUE: half - VALUE: float - VALUE: int32 shader_variants: - - NAME: arange + - NAME: arange_buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl new file mode 100644 index 00000000000..04631e0ef06 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl @@ -0,0 +1,52 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#version 450 core + +${define_required_extensions("texture3d", DTYPE)} + +#define PRECISION ${PRECISION} + +#define VEC4_T ${texel_load_type(DTYPE, "texture3d")} + +${define_active_storage_type("texture3d")} + +layout(std430) buffer; + +#include "indexing.glslh" + +${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(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const ivec3 out_pos = ivec3(gl_GlobalInvocationID); + + if (out_of_bounds(out_pos, outp)) { + return; + } + + TensorIndex4D out_tidx = texture_pos_to_tensor4d_idx_simple(outp, out_pos); + + // 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); + int limit = min( + 4, outp.sizes[outp.packed_dim] - out_tidx.data[outp.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; + out_tidx.data[outp.packed_dim]++; + } + + imageStore(t_out, out_pos, outtex); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.yaml new file mode 100644 index 00000000000..2cd9255d754 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.yaml @@ -0,0 +1,16 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +arange_texture: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + - VALUE: int32 + shader_variants: + - NAME: arange_texture3d diff --git a/backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl new file mode 100644 index 00000000000..729baec324d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl @@ -0,0 +1,36 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#version 450 core + +${define_required_extensions("buffer", DTYPE)} + +#define PRECISION ${PRECISION} + +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} + +layout(std430) buffer; + +#include "indexing.glslh" + +${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} +${layout_declare_ubo(B, "BufferMetadata", "outp")} +${layout_declare_ubo(B, "float", "fill_value")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const uint out_bufi = gl_GlobalInvocationID.x; + if (out_of_bounds(out_bufi, outp)) { + return; + } + + t_out[out_bufi] = T(fill_value); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/full_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/full_buffer.yaml new file mode 100644 index 00000000000..7b1af0ab9c6 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/full_buffer.yaml @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +full_buffer: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + - VALUE: int32 + - VALUE: uint8 + shader_variants: + - NAME: full_buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/full.glsl b/backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl similarity index 59% rename from backends/vulkan/runtime/graph/ops/glsl/full.glsl rename to backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl index 81f1f182cdf..eded1d892c1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/full.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl @@ -8,35 +8,36 @@ #version 450 core -#define PRECISION ${PRECISION} +${define_required_extensions("texture3d", DTYPE)} -#define VEC4_T ${texel_type(DTYPE)} +#define PRECISION ${PRECISION} -#define POS ${get_pos[NDIM]("pos")} +#define VEC4_T ${texel_load_type(DTYPE, "texture3d")} -#include "indexing_utils.h" +${define_active_storage_type("texture3d")} layout(std430) buffer; -${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} -${layout_declare_ubo(B, "ivec4", "sizes")} +#include "indexing.glslh" + +${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} +${layout_declare_ubo(B, "TextureMetadata", "outp")} ${layout_declare_ubo(B, "float", "fill_value")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -layout(constant_id = 3) const int packed_dim = C_DIM; - void main() { const ivec3 pos = ivec3(gl_GlobalInvocationID); - const ivec4 idx = to_tensor_idx(pos, sizes, packed_dim); - if (any(greaterThanEqual(idx, sizes))) { + if (out_of_bounds(pos, outp)) { return; } VEC4_T outtex = VEC4_T(fill_value); - const int packed_dim_size = sizes[packed_dim]; - int packed_idx = idx[packed_dim]; + + TensorIndex4D tidx = texture_pos_to_tensor4d_idx_simple(outp, pos); + const int packed_dim_size = outp.sizes[outp.packed_dim]; + int packed_idx = tidx.data[outp.packed_dim]; if (packed_idx + 3 >= packed_dim_size) { ivec4 packed_ind = ivec4(packed_idx) + ivec4(0, 1, 2, 3); @@ -44,5 +45,5 @@ void main() { outtex = outtex * valid_idx; } - imageStore(t_out, POS, outtex); + imageStore(t_out, pos, outtex); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/full.yaml b/backends/vulkan/runtime/graph/ops/glsl/full_texture.yaml similarity index 88% rename from backends/vulkan/runtime/graph/ops/glsl/full.yaml rename to backends/vulkan/runtime/graph/ops/glsl/full_texture.yaml index 5d7a983cae3..80e87dde5fd 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/full.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/full_texture.yaml @@ -4,11 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -full: +full_texture: parameter_names_with_default_values: - NDIM: 3 DTYPE: float - PACKING: C_packed STORAGE: texture3d generate_variant_forall: DTYPE: @@ -17,4 +15,4 @@ full: - VALUE: int32 - VALUE: uint8 shader_variants: - - NAME: full + - NAME: full_texture3d diff --git a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp index 3171fbeb488..ed4f7313475 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp @@ -85,6 +85,7 @@ void add_arange_node( 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)); graph.execute_nodes().emplace_back(new DynamicDispatchNode( @@ -95,7 +96,7 @@ void add_arange_node( // Inputs and Outputs {{out, vkapi::kWrite}}, // Shader params buffers - {graph.sizes_ubo(out), + {graph.meta_ubo(out), graph.create_params_buffer(start_val), graph.create_params_buffer(step_val)}, // Push Constants diff --git a/backends/vulkan/runtime/graph/ops/impl/Full.cpp b/backends/vulkan/runtime/graph/ops/impl/Full.cpp index fe2676e91e0..42f1f050e37 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Full.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Full.cpp @@ -25,7 +25,7 @@ void resize_full_node( if (graph->val_is_tensor(extra_args.at(0))) { out_sizes = graph->sizes_of(extra_args.at(0)); } else { - out_sizes = *graph->get_int_list(extra_args.at(0)); + out_sizes = graph->extract_int_or_symint_list(extra_args.at(0)); } graph->virtual_resize(out, out_sizes); @@ -41,6 +41,7 @@ void add_full_node( std::string kernel_name("full"); kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); graph.execute_nodes().emplace_back(new DynamicDispatchNode( @@ -51,11 +52,11 @@ void add_full_node( // Inputs and Outputs {{out, vkapi::kWrite}}, // Shader params buffers - {graph.sizes_ubo(out), graph.create_params_buffer(fill_value_val)}, + {graph.meta_ubo(out), graph.create_params_buffer(fill_value_val)}, // Push Constants {}, // Specialization Constants - {graph.packed_dim_of(out)}, + {}, // Resize Args {size_or_in}, // Resizing Logic diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py index 3b2d6b8c48e..1c4356cab4f 100644 --- a/backends/vulkan/test/op_tests/cases.py +++ b/backends/vulkan/test/op_tests/cases.py @@ -802,6 +802,7 @@ def get_full_inputs(): ([L, M, M1, M2], 2.72), ] ) + test_suite.storage_types = ["utils::kTexture3D", "utils::kBuffer"] return test_suite @@ -836,6 +837,7 @@ def get_ones_inputs(): ([L, M, M1, M2]), ] ) + test_suite.storage_types = ["utils::kTexture3D", "utils::kBuffer"] return test_suite @@ -1849,8 +1851,13 @@ def get_arange_inputs(): ) test_suite.layouts = [ + "utils::kWidthPacked", "utils::kChannelsPacked", ] + test_suite.storage_types = [ + "utils::kTexture3D", + "utils::kBuffer", + ] return test_suite