[ET-VK][ops] Add buffer unfold support

Implement `aten.unfold_copy.default` for Vulkan buffer tensors, including high-rank indexing and dynamic resize support. Add focused coverage for general, high-rank, and SceneX audio shapes. Authored with Codex.

Differential Revision: [D114382687](https://our.internmc.facebook.com/intern/diff/D114382687/)


ghstack-source-id: 409188517
Pull-Request: pytorch#21521

Author: SS-JIA

backends/vulkan/op_registry.py

@@ -226,6 +226,7 @@ def register_copy_op():
         exir_ops.edge.aten.tanh.default,
         exir_ops.edge.aten.round.default,
         exir_ops.edge.aten.leaky_relu.default,
+        exir_ops.edge.aten.log10.default,
     ]
 )
 def register_unaryop_cpp_ops():
@@ -1271,6 +1272,21 @@ def register_alias_copy():
     )
 
 
+# =============================================================================
+# Unfold.cpp
+# =============================================================================
+
+
+@update_features(exir_ops.edge.aten.unfold_copy.default)
+def register_unfold_copy():
+    return OpFeatures(
+        inputs_storage=utils.ANY_BUFFER,
+        inputs_dtypes=utils.FP_T,
+        supports_resize=True,
+        supports_highdim=True,
+    )
+
+
 # =============================================================================
 # Gather.cpp
 # =============================================================================

backends/vulkan/runtime/graph/ops/glsl/unfold_copy_buffer.glsl

@@ -0,0 +1,51 @@
+/*
+ * 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_tensor(B, "r", "t_in", DTYPE, "buffer")}
+
+${layout_declare_ubo(B, "BufferMetadata", "out_meta")}
+${layout_declare_ubo(B, "BufferMetadata", "in_meta")}
+
+${layout_declare_spec_const(C, "int", "unfold_dim", "0")}
+${layout_declare_spec_const(C, "int", "step", "1")}
+
+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, out_meta)) {
+    return;
+  }
+
+  const TensorIndex out_tidx = linear_idx_to_tensor_idx(out_meta, out_bufi);
+  TensorIndex in_tidx;
+  initialize(in_tidx);
+
+  for (int d = 0; d < int_ndim(in_meta); ++d) {
+    in_tidx.data[div_4(d)][mod_4(d)] = idx_at(out_tidx, d + 1);
+  }
+  in_tidx.data[div_4(unfold_dim)][mod_4(unfold_dim)] =
+      idx_at(out_tidx, unfold_dim + 1) * step + idx_at(out_tidx, 0);
+
+  const uint in_bufi = tensor_idx_to_linear_idx(in_meta, in_tidx);
+  t_out[out_bufi] = t_in[in_bufi];
+}

backends/vulkan/runtime/graph/ops/glsl/unfold_copy_buffer.yaml

@@ -0,0 +1,15 @@
+# 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.
+
+unfold_copy_buffer:
+  parameter_names_with_default_values:
+    DTYPE: float
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: unfold_copy_buffer

backends/vulkan/runtime/graph/ops/impl/Unfold.cpp

@@ -0,0 +1,109 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+int64_t normalize_unfold_dim(const int64_t dim, const int64_t ndim) {
+  const int64_t normalized_dim = dim < 0 ? dim + ndim : dim;
+  VK_CHECK_COND(
+      normalized_dim >= 0 && normalized_dim < ndim,
+      "unfold_copy: dimension must be in range [-input.dim(), input.dim())");
+  return normalized_dim;
+}
+
+std::vector<int64_t> get_unfold_sizes(
+    ComputeGraph& graph,
+    const ValueRef input,
+    const ValueRef dim_ref,
+    const ValueRef size_ref,
+    const ValueRef step_ref) {
+  const std::vector<int64_t> input_sizes = graph.sizes_of(input);
+  VK_CHECK_COND(!input_sizes.empty(), "unfold_copy: input must have rank >= 1");
+
+  const int64_t dim = normalize_unfold_dim(
+      graph.extract_scalar<int64_t>(dim_ref), input_sizes.size());
+  const int64_t size = graph.extract_scalar<int64_t>(size_ref);
+  const int64_t step = graph.extract_scalar<int64_t>(step_ref);
+
+  VK_CHECK_COND(size > 0, "unfold_copy: size must be greater than zero");
+  VK_CHECK_COND(step > 0, "unfold_copy: step must be greater than zero");
+  VK_CHECK_COND(
+      size <= input_sizes.at(dim),
+      "unfold_copy: size must not exceed selected input dimension");
+
+  std::vector<int64_t> output_sizes = input_sizes;
+  output_sizes.at(dim) = (input_sizes.at(dim) - size) / step + 1;
+  output_sizes.push_back(size);
+  return output_sizes;
+}
+
+void resize_unfold_copy_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  const ValueRef output = args.at(0).refs.at(0);
+  const ValueRef input = args.at(1).refs.at(0);
+  graph->virtual_resize(
+      output,
+      get_unfold_sizes(
+          *graph,
+          input,
+          resize_args.at(0),
+          resize_args.at(1),
+          resize_args.at(2)));
+}
+
+void add_unfold_copy_node(
+    ComputeGraph& graph,
+    const ValueRef input,
+    const ValueRef dim_ref,
+    const ValueRef size_ref,
+    const ValueRef step_ref,
+    const ValueRef output) {
+  const int64_t input_ndim = graph.dim_of(input);
+  const int64_t dim =
+      normalize_unfold_dim(graph.extract_scalar<int64_t>(dim_ref), input_ndim);
+  const int64_t dim_whcn = input_ndim - dim - 1;
+  const int64_t step = graph.extract_scalar<int64_t>(step_ref);
+
+  get_unfold_sizes(graph, input, dim_ref, size_ref, step_ref);
+
+  std::string kernel_name = "unfold_copy_buffer";
+  add_dtype_suffix(kernel_name, graph.dtype_of(output));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      default_pick_global_wg_size,
+      default_pick_local_wg_size,
+      {{output, vkapi::kWrite}, {input, vkapi::kRead}},
+      {graph.meta_ubo(output), graph.meta_ubo(input)},
+      {},
+      {
+          utils::safe_downcast<int32_t>(dim_whcn),
+          utils::safe_downcast<int32_t>(step),
+      },
+      {dim_ref, size_ref, step_ref},
+      resize_unfold_copy_node));
+}
+
+void unfold_copy(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  add_unfold_copy_node(
+      graph, args.at(0), args.at(1), args.at(2), args.at(3), args.at(4));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(aten.unfold_copy.default, unfold_copy);
+}
+
+} // namespace vkcompute

backends/vulkan/test/op_tests/cases.py

@@ -1834,6 +1834,54 @@ def get_unary_ops_inputs():
     return test_suite
 
 
+@register_test_suite("aten.unfold_copy.default")
+def get_unfold_copy_inputs():
+    Test = namedtuple("UnfoldCopy", ["self", "dimension", "size", "step"])
+
+    test_suite = VkTestSuite(
+        [
+            Test(self=(11,), dimension=0, size=5, step=2),
+            Test(self=(3, 11), dimension=-1, size=5, step=2),
+            Test(self=(3, 7, 11), dimension=1, size=3, step=2),
+            Test(self=(5, 3, 7, 11), dimension=0, size=3, step=2),
+        ]
+    )
+    test_suite.storage_types = ["utils::kBuffer"]
+    test_suite.layouts = [
+        "utils::kWidthPacked",
+        "utils::kChannelsPacked",
+    ]
+    test_suite.data_gen = "make_seq_tensor"
+    test_suite.test_name_suffix = "buffer"
+
+    highdim_test_suite = VkTestSuite(
+        [
+            Test(
+                self=(2, 2, 2, 3, 2, 2, 2),
+                dimension=-4,
+                size=2,
+                step=1,
+            ),
+        ]
+    )
+    highdim_test_suite.storage_types = ["utils::kBuffer"]
+    highdim_test_suite.layouts = ["utils::kChannelsPacked"]
+    highdim_test_suite.data_gen = "make_seq_tensor"
+    highdim_test_suite.test_name_suffix = "highdim_buffer"
+
+    scenex_test_suite = VkTestSuite(
+        [
+            Test(self=(30, 15760), dimension=-1, size=400, step=160),
+        ]
+    )
+    scenex_test_suite.storage_types = ["utils::kBuffer"]
+    scenex_test_suite.layouts = ["utils::kWidthPacked"]
+    scenex_test_suite.dtypes = ["at::kFloat"]
+    scenex_test_suite.test_name_suffix = "scenex_buffer"
+
+    return [test_suite, highdim_test_suite, scenex_test_suite]
+
+
 # separate test suite from unary_ops for learning purposes
 @register_test_suite("aten.tan.default")
 def get_tan_inputs():