[ET-VK] Add ANY_STORAGE support to arange and full

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,
     )
 

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);
+}

backends/vulkan/runtime/graph/ops/glsl/arange.yaml → 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

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);
+}

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

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);
+}

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

backends/vulkan/runtime/graph/ops/glsl/full.glsl → backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl

@@ -8,41 +8,42 @@
 
 #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);
     VEC4_T valid_idx = VEC4_T(lessThan(packed_ind, ivec4(packed_dim_size)));
     outtex = outtex * valid_idx;
   }
 
-  imageStore(t_out, POS, outtex);
+  imageStore(t_out, pos, outtex);
 }

backends/vulkan/runtime/graph/ops/glsl/full.yaml → 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

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

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

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