Copy slice-scatter inputs into channels-packed textures

backends/native/runtime/vulkan/VulkanEngineExecutable.cpp

@@ -1749,6 +1749,23 @@
       }
     }
 
+    // ET-VK's slice_scatter indexes self, src and out as standard
+    // channels-packed textures.
+    const bool slice_scatter = key == "aten.slice_scatter.default";
+    if (slice_scatter) {
+      const uint32_t limit =
+          graph_->context()->adapter_ptr()->max_texture3d_dim();
+      std::vector<ValueId> ids = n.input_value_ids();
+      for_each_output_value(n, [&](ValueId id) { ids.push_back(id); });
+      for (const ValueId id : ids) {
+        if (g().value(id).is_tensor() &&
+            exceeds_texture_limit(meta_of(id), utils::kChannelsPacked, limit)) {
+          throw std::runtime_error(
+              "vulkan: slice_scatter exceeds the texture dimension limit");
+        }
+      }
+    }
+
     std::vector<VkRef> args;
     args.reserve(n.inputs.size() + 1);
     for (const NamedArgument& na : n.inputs) {
@@ -1786,6 +1803,16 @@
             storage_at(value_id),
             utils::kChannelsPacked);
       }
+      if (slice_scatter && (na.name == "self" || na.name == "src") &&
+          na.arg.kind() == ArgKind::Tensor && valid(na.arg.as_tensor().id) &&
+          !graph_->val_is_tref(a) &&
+          !graph_->is_standard_channels_packed_texture_tensor(a)) {
+        a = repack(
+            a,
+            g().value(na.arg.as_tensor().id),
+            utils::kTexture3D,
+            utils::kChannelsPacked);
+      }
       // ET-VK's pad picks its shader by the output's layout and storage
       // alone, which later layout steps can leave unlike its input's.
       if (key == "aten.constant_pad_nd.default" && na.name == "self" &&
@@ -1828,6 +1855,18 @@
     std::vector<VkRef> outs;
     for_each_output_value(
         n, [&](ValueId output_id) { outs.push_back(vref_at(output_id)); });
+    if (slice_scatter && outs.size() == 1 &&
+        !graph_->is_standard_channels_packed_texture_tensor(outs[0])) {
+      const VkRef tmp = graph_->add_tensor(
+          graph_->sizes_of(outs[0]),
+          graph_->dtype_of(outs[0]),
+          utils::kTexture3D,
+          utils::kChannelsPacked);
+      args.push_back(tmp);
+      VK_GET_OP_FN(key)(*graph_, args);
+      VK_GET_OP_FN("aten.clone.default")(*graph_, {tmp, none_ref_, outs[0]});
+      return;
+    }
     // ET-VK's embedding and convolution only write channels-packed outputs.
     if ((key == "aten.embedding.default" ||
          key == "aten.convolution.default") &&

backends/native/test/runtime/vulkan/test_vulkan_engine.cpp

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -1370,6 +1370,106 @@
   return tensors;
 }
 
+// result = index(slice_scatter(index(input, batch), index(src, batch), 1, 1, 5,
+// 2), batch) over a [2, 6] input and [2, 2] src. aten.index keeps its operands
+// in buffer storage, so slice_scatter reads and writes buffer-planned tensors.
+std::vector<uint8_t> make_buffered_slice_scatter_program() {
+  flatbuffers::FlatBufferBuilder builder;
+  const auto meta = create_tensor_meta(builder, fbs::ScalarType::FLOAT, {2, 6});
+  const std::vector<flatbuffers::Offset<fbs::TensorValue>> tensor_values = {
+      fbs::CreateTensorValueDirect(builder, "input", meta),
+      fbs::CreateTensorValueDirect(
+          builder,
+          "src",
+          create_tensor_meta(builder, fbs::ScalarType::FLOAT, {2, 2})),
+      fbs::CreateTensorValueDirect(
+          builder,
+          "batch",
+          create_tensor_meta(builder, fbs::ScalarType::LONG, {2})),
+      fbs::CreateTensorValueDirect(builder, "gathered", meta),
+      fbs::CreateTensorValueDirect(
+          builder,
+          "gathered_src",
+          create_tensor_meta(builder, fbs::ScalarType::FLOAT, {2, 2})),
+      fbs::CreateTensorValueDirect(builder, "scattered", meta),
+      fbs::CreateTensorValueDirect(builder, "result", meta),
+  };
+  std::vector<flatbuffers::Offset<fbs::Node>> nodes;
+  for (const char* name : {"input", "src", "batch"}) {
+    const std::vector<flatbuffers::Offset<fbs::Output>> outputs = {
+        fbs::CreateOutputDirect(builder, name)};
+    nodes.push_back(
+        fbs::CreateNodeDirect(
+            builder, name, fbs::OpKind::PLACEHOLDER, "", nullptr, &outputs));
+  }
+  const auto add_call =
+      [&](const char* name,
+          const char* target,
+          std::vector<flatbuffers::Offset<fbs::NamedArgument>> inputs) {
+        const std::vector<flatbuffers::Offset<fbs::Output>> outputs = {
+            fbs::CreateOutputDirect(builder, name)};
+        nodes.push_back(
+            fbs::CreateNodeDirect(
+                builder,
+                name,
+                fbs::OpKind::CALL_FUNCTION,
+                target,
+                &inputs,
+                &outputs));
+      };
+  const auto gather_inputs = [&](const char* input) {
+    const auto indices = fbs::CreateOptionalTensorListArg(
+        builder,
+        create_strings(builder, {"batch"}),
+        builder.CreateVector(std::vector<uint8_t>{1}));
+    return std::vector<flatbuffers::Offset<fbs::NamedArgument>>{
+        fbs::CreateNamedArgumentDirect(
+            builder, "self", create_tensor_arg(builder, input)),
+        fbs::CreateNamedArgumentDirect(
+            builder,
+            "indices",
+            fbs::CreateArgument(
+                builder,
+                fbs::ArgumentValue::OptionalTensorListArg,
+                indices.Union())),
+    };
+  };
+  add_call("gathered", "torch.ops.aten.index.Tensor", gather_inputs("input"));
+  add_call("gathered_src", "torch.ops.aten.index.Tensor", gather_inputs("src"));
+  add_call(
+      "scattered",
+      "torch.ops.aten.slice_scatter.default",
+      {
+          fbs::CreateNamedArgumentDirect(
+              builder, "self", create_tensor_arg(builder, "gathered")),
+          fbs::CreateNamedArgumentDirect(
+              builder, "src", create_tensor_arg(builder, "gathered_src")),
+          fbs::CreateNamedArgumentDirect(
+              builder, "dim", create_int_arg(builder, 1)),
+          fbs::CreateNamedArgumentDirect(
+              builder, "start", create_int_arg(builder, 1)),
+          fbs::CreateNamedArgumentDirect(
+              builder, "end", create_int_arg(builder, 5)),
+          fbs::CreateNamedArgumentDirect(
+              builder, "step", create_int_arg(builder, 2)),
+      });
+  add_call("result", "torch.ops.aten.index.Tensor", gather_inputs("scattered"));
+  const std::vector<flatbuffers::Offset<fbs::NamedArgument>> output_inputs = {
+      fbs::CreateNamedArgumentDirect(
+          builder, "", create_tensor_arg(builder, "result"))};
+  nodes.push_back(
+      fbs::CreateNodeDirect(
+          builder, "output", fbs::OpKind::OUTPUT, "", &output_inputs));
+  const auto graph = fbs::CreateGraph(
+      builder,
+      builder.CreateVector(nodes),
+      create_strings(builder, {"input", "src", "batch"}),
+      create_strings(builder, {"result"}),
+      builder.CreateVector(tensor_values));
+  return finish_program(
+      builder, graph, {fbs::CreateOutputSpecDirect(builder, "result")});
+}
+
 constexpr int64_t kHfTokens = 4;
 constexpr int64_t kHfHeads = 4;
 constexpr int64_t kHfKvHeads = 2;
@@ -2182,6 +2282,32 @@
   EXPECT_EQ(output, (std::vector<float>{4.0f, 1.0f}));
 }
 
+TEST(VulkanEngineTest, ExecutesSliceScatterOnBufferPlannedTensors) {
+  std::vector<float> input(12);
+  for (size_t i = 0; i < input.size(); ++i) {
+    input[i] = static_cast<float>(i + 1);
+  }
+  const std::vector<float> src = {-1.0f, -2.0f, -3.0f, -4.0f};
+  const std::vector<int64_t> batch = {0, 1};
+  CompiledProgram compiled =
+      compile_program(make_buffered_slice_scatter_program());
+  compiled.executable->set_input(
+      0, input.data(), input.size(), ScalarType::Float);
+  compiled.executable->set_input(1, src.data(), src.size(), ScalarType::Float);
+  compiled.executable->set_input(
+      2, batch.data(), batch.size(), ScalarType::Long);
+  compiled.executable->execute();
+  std::vector<float> output(12);
+  compiled.executable->get_output(
+      0, output.data(), output.size(), ScalarType::Float);
+  std::vector<float> expected = input;
+  for (size_t row = 0; row < 2; ++row) {
+    expected[row * 6 + 1] = src[row * 2];
+    expected[row * 6 + 3] = src[row * 2 + 1];
+  }
+  EXPECT_EQ(output, expected);
+}
+
 TEST(VulkanEngineTest, RejectsBufferViewResizeThatDoesNotDivide) {
   CompiledProgram compiled =
       compile_program(make_dynamic_buffer_view_index_program());

examples/native/vision/BUCK

@@ -12,7 +12,6 @@ fbcode_target(
     deps = [
         "fbsource//third-party/pypi/pillow:pillow",
         "//caffe2:torch",
-        "//executorch/examples/models:models",
         "//executorch/exir/native:native",
         "//pytorch/vision:torchvision",  # @manual
     ],

examples/native/vision/export.py

@@ -17,15 +17,17 @@
 from typing import TypedDict
 
 import torch
-from executorch.examples.models import MODEL_NAME_TO_MODEL
-from executorch.examples.models.model_factory import EagerModelFactory
 from executorch.exir.native import to_native
 from PIL import Image
-from torchvision.models import MobileNet_V2_Weights, ResNet50_Weights  # @manual
+from torchvision import models  # @manual
 
 logger: logging.Logger = logging.getLogger(__name__)
 
-_MODEL_NAMES = ("mv2", "resnet50")
+_MODEL_ALIASES = {"mv2": "mobilenet_v2", "vit": "vit_b_16"}
+_WEIGHTS = {
+    name: models.get_model_weights(name)["DEFAULT"]
+    for name in models.list_models(module=models)
+}
 
 
 class _Expected(TypedDict):
@@ -41,22 +43,14 @@ def _default_out_dir() -> Path:
 
 
 def _load_model(name: str) -> torch.nn.Module:
-    model, _, _, _ = EagerModelFactory.create_model(*MODEL_NAME_TO_MODEL[name])
-    return model.eval()
+    return models.get_model(name, weights=_WEIGHTS[name]).eval()
 
 
 def _load_image(name: str, image_path: Path) -> tuple[torch.Tensor, list[str]]:
+    weights = _WEIGHTS[name]
     with Image.open(image_path) as image:
-        rgb_image = image.convert("RGB")
-        if name == "mv2":
-            weights = MobileNet_V2_Weights.DEFAULT
-            image_tensor = weights.transforms()(rgb_image)
-            categories = list(weights.meta["categories"])
-        else:
-            weights = ResNet50_Weights.IMAGENET1K_V1
-            image_tensor = weights.transforms()(rgb_image)
-            categories = list(weights.meta["categories"])
-    return image_tensor.unsqueeze(0).contiguous(), categories
+        image_tensor = weights.transforms()(image.convert("RGB"))
+    return image_tensor.unsqueeze(0).contiguous(), list(weights.meta["categories"])
 
 
 def _expected(
@@ -113,7 +107,7 @@ def _summary(
 def main() -> None:
     logging.basicConfig(level=logging.INFO)
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model", choices=_MODEL_NAMES, default="mv2")
+    parser.add_argument("--model", choices=[*_WEIGHTS, *_MODEL_ALIASES], default="mv2")
     parser.add_argument("--out-dir", type=Path, default=None)
     parser.add_argument("--image", type=Path, required=True)
     args = parser.parse_args()
@@ -129,8 +123,9 @@ def main() -> None:
         "summary": output_dir / f"{args.model}.summary.txt",
     }
 
-    model = _load_model(args.model)
-    example_input, categories = _load_image(args.model, args.image)
+    model_name = _MODEL_ALIASES.get(args.model, args.model)
+    model = _load_model(model_name)
+    example_input, categories = _load_image(model_name, args.image)
     logits, expected = _expected(model, example_input, categories)
     native_program = to_native(torch.export.export(model, (example_input,)))