Prepack stored inputs of dtype conversions

backends/native/runtime/vulkan/passes/InsertPrepack.cpp

@@ -19,7 +19,8 @@ namespace ptn::vulkan {
 namespace {
 
 bool needs_prepack(const Node& node) {
-  constexpr std::array<std::string_view, 4> kOps{
+  constexpr std::array<std::string_view, 5> kOps{
+      "torch.ops.aten._to_copy.default",
       "torch.ops.aten.add.Tensor",
       "torch.ops.aten.index.Tensor",
       "torch.ops.aten.mul.Tensor",

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

@@ -98,6 +98,15 @@ TEST(InsertPrepackTest, PrepacksConstantSliceInput) {
       method.graph.node(2).input_value_ids(), std::vector<ValueId>({0, 4}));
 }
 
+TEST(InsertPrepackTest, PrepacksConstantToCopyInput) {
+  Method method = make_method();
+  method.graph.node(2).target = "torch.ops.aten._to_copy.default";
+
+  EXPECT_EQ(insert_prepack_nodes(method), 1);
+  EXPECT_EQ(
+      method.graph.node(2).input_value_ids(), std::vector<ValueId>({0, 4}));
+}
+
 TEST(InsertPrepackTest, LeavesSelfPrepackingOpUntouched) {
   Method method = make_method();
   method.graph.node(2).target = "torch.ops.aten.linear.default";

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

@@ -1370,6 +1370,109 @@ std::vector<uint8_t> make_quantized_input_tensors() {
   return tensors;
 }
 
+// sum = input + _to_copy(offsets, float), where `offsets` is a stored int64
+// constant.
+std::vector<uint8_t> make_stored_long_to_copy_program() {
+  flatbuffers::FlatBufferBuilder builder;
+  const auto float_meta =
+      create_tensor_meta(builder, fbs::ScalarType::FLOAT, {4});
+  const auto long_meta =
+      create_tensor_meta(builder, fbs::ScalarType::LONG, {4});
+  const std::vector<flatbuffers::Offset<fbs::TensorValue>> tensor_values = {
+      fbs::CreateTensorValueDirect(builder, "input", float_meta),
+      fbs::CreateTensorValueDirect(builder, "offsets", long_meta),
+      fbs::CreateTensorValueDirect(builder, "converted", float_meta),
+      fbs::CreateTensorValueDirect(builder, "sum", float_meta),
+  };
+  const std::vector<flatbuffers::Offset<fbs::Output>> input_outputs = {
+      fbs::CreateOutputDirect(builder, "input")};
+  const std::vector<flatbuffers::Offset<fbs::Output>> offsets_outputs = {
+      fbs::CreateOutputDirect(builder, "offsets")};
+  const std::vector<flatbuffers::Offset<fbs::Output>> converted_outputs = {
+      fbs::CreateOutputDirect(builder, "converted")};
+  const std::vector<flatbuffers::Offset<fbs::Output>> sum_outputs = {
+      fbs::CreateOutputDirect(builder, "sum")};
+  const auto dtype = fbs::CreateScalarTypeArg(builder, fbs::ScalarType::FLOAT);
+  const std::vector<flatbuffers::Offset<fbs::NamedArgument>> to_copy_inputs = {
+      fbs::CreateNamedArgumentDirect(
+          builder, "self", create_tensor_arg(builder, "offsets")),
+      fbs::CreateNamedArgumentDirect(
+          builder,
+          "dtype",
+          fbs::CreateArgument(
+              builder, fbs::ArgumentValue::ScalarTypeArg, dtype.Union())),
+      fbs::CreateNamedArgumentDirect(
+          builder, "layout", create_none_arg(builder)),
+      fbs::CreateNamedArgumentDirect(
+          builder, "device", create_none_arg(builder)),
+      fbs::CreateNamedArgumentDirect(
+          builder, "pin_memory", create_none_arg(builder)),
+      fbs::CreateNamedArgumentDirect(
+          builder, "non_blocking", create_bool_arg(builder, false)),
+      fbs::CreateNamedArgumentDirect(
+          builder, "memory_format", create_none_arg(builder)),
+  };
+  const std::vector<flatbuffers::Offset<fbs::NamedArgument>> add_inputs = {
+      fbs::CreateNamedArgumentDirect(
+          builder, "self", create_tensor_arg(builder, "input")),
+      fbs::CreateNamedArgumentDirect(
+          builder, "other", create_tensor_arg(builder, "converted")),
+      fbs::CreateNamedArgumentDirect(
+          builder, "alpha", create_int_arg(builder, 1)),
+  };
+  const std::vector<flatbuffers::Offset<fbs::NamedArgument>> output_inputs = {
+      fbs::CreateNamedArgumentDirect(
+          builder, "", create_tensor_arg(builder, "sum"))};
+  const std::vector<flatbuffers::Offset<fbs::Node>> nodes = {
+      fbs::CreateNodeDirect(
+          builder,
+          "input",
+          fbs::OpKind::PLACEHOLDER,
+          "",
+          nullptr,
+          &input_outputs),
+      fbs::CreateNodeDirect(
+          builder,
+          "offsets",
+          fbs::OpKind::PLACEHOLDER,
+          "",
+          nullptr,
+          &offsets_outputs),
+      fbs::CreateNodeDirect(
+          builder,
+          "converted",
+          fbs::OpKind::CALL_FUNCTION,
+          "torch.ops.aten._to_copy.default",
+          &to_copy_inputs,
+          &converted_outputs),
+      fbs::CreateNodeDirect(
+          builder,
+          "sum",
+          fbs::OpKind::CALL_FUNCTION,
+          "torch.ops.aten.add.Tensor",
+          &add_inputs,
+          &sum_outputs),
+      fbs::CreateNodeDirect(
+          builder, "output", fbs::OpKind::OUTPUT, "", &output_inputs),
+  };
+  const auto graph = fbs::CreateGraph(
+      builder,
+      builder.CreateVector(nodes),
+      create_strings(builder, {"input"}),
+      create_strings(builder, {"sum"}),
+      builder.CreateVector(tensor_values));
+  const std::vector<flatbuffers::Offset<fbs::NamedTensorRef>> constants = {
+      fbs::CreateNamedTensorRefDirect(
+          builder,
+          "offsets",
+          "offsets",
+          long_meta,
+          fbs::InputKind::CONSTANT_TENSOR),
+  };
+  return finish_program(
+      builder, graph, {fbs::CreateOutputSpecDirect(builder, "sum")}, constants);
+}
+
 constexpr int64_t kHfTokens = 4;
 constexpr int64_t kHfHeads = 4;
 constexpr int64_t kHfKvHeads = 2;
@@ -2182,6 +2285,27 @@ TEST(VulkanEngineTest, ResizesBufferViewWithItsSource) {
   EXPECT_EQ(output, (std::vector<float>{4.0f, 1.0f}));
 }
 
+TEST(VulkanEngineTest, ConvertsStoredLongConstant) {
+  const std::string header =
+      R"({"offsets":{"dtype":"I64","shape":[4],"data_offsets":[0,32]}})";
+  std::vector<uint8_t> tensors = testing::make_safetensors(header, 32);
+  const std::vector<int64_t> offsets = {3, -2, 7, 100};
+  std::memcpy(
+      tensors.data() + sizeof(uint64_t) + header.size(),
+      offsets.data(),
+      offsets.size() * sizeof(int64_t));
+  CompiledProgram compiled =
+      compile_program(make_stored_long_to_copy_program(), std::move(tensors));
+  const std::vector<float> input = {0.5f, 1.5f, 2.5f, 3.5f};
+  compiled.executable->set_input(
+      0, input.data(), input.size(), ScalarType::Float);
+  compiled.executable->execute();
+  std::vector<float> output(4);
+  compiled.executable->get_output(
+      0, output.data(), output.size(), ScalarType::Float);
+  EXPECT_EQ(output, (std::vector<float>{3.5f, -0.5f, 9.5f, 103.5f}));
+}
+
 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,)))