diff --git a/backends/native/runtime/vulkan/passes/InsertPrepack.cpp b/backends/native/runtime/vulkan/passes/InsertPrepack.cpp index 776cffea3a4..10ce15d3c8e 100644 --- a/backends/native/runtime/vulkan/passes/InsertPrepack.cpp +++ b/backends/native/runtime/vulkan/passes/InsertPrepack.cpp @@ -19,7 +19,8 @@ namespace ptn::vulkan { namespace { bool needs_prepack(const Node& node) { - constexpr std::array kOps{ + constexpr std::array kOps{ + "torch.ops.aten._to_copy.default", "torch.ops.aten.add.Tensor", "torch.ops.aten.index.Tensor", "torch.ops.aten.mul.Tensor", diff --git a/backends/native/test/runtime/vulkan/test_insert_prepack.cpp b/backends/native/test/runtime/vulkan/test_insert_prepack.cpp index 168d2955b58..8d92fafc886 100644 --- a/backends/native/test/runtime/vulkan/test_insert_prepack.cpp +++ b/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({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({0, 4})); +} + TEST(InsertPrepackTest, LeavesSelfPrepackingOpUntouched) { Method method = make_method(); method.graph.node(2).target = "torch.ops.aten.linear.default"; diff --git a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp index 78543924005..516ebf074ff 100644 --- a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp +++ b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp @@ -1370,6 +1370,109 @@ std::vector make_quantized_input_tensors() { return tensors; } +// sum = input + _to_copy(offsets, float), where `offsets` is a stored int64 +// constant. +std::vector 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> 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> input_outputs = { + fbs::CreateOutputDirect(builder, "input")}; + const std::vector> offsets_outputs = { + fbs::CreateOutputDirect(builder, "offsets")}; + const std::vector> converted_outputs = { + fbs::CreateOutputDirect(builder, "converted")}; + const std::vector> sum_outputs = { + fbs::CreateOutputDirect(builder, "sum")}; + const auto dtype = fbs::CreateScalarTypeArg(builder, fbs::ScalarType::FLOAT); + const std::vector> 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> 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> output_inputs = { + fbs::CreateNamedArgumentDirect( + builder, "", create_tensor_arg(builder, "sum"))}; + const std::vector> 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> 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{4.0f, 1.0f})); } +TEST(VulkanEngineTest, ConvertsStoredLongConstant) { + const std::string header = + R"({"offsets":{"dtype":"I64","shape":[4],"data_offsets":[0,32]}})"; + std::vector tensors = testing::make_safetensors(header, 32); + const std::vector 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 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 output(4); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::Float); + EXPECT_EQ(output, (std::vector{3.5f, -0.5f, 9.5f, 103.5f})); +} + TEST(VulkanEngineTest, RejectsBufferViewResizeThatDoesNotDivide) { CompiledProgram compiled = compile_program(make_dynamic_buffer_view_index_program()); diff --git a/examples/native/vision/BUCK b/examples/native/vision/BUCK index fb27e3a34af..008e7cc7c7d 100644 --- a/examples/native/vision/BUCK +++ b/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 ], diff --git a/examples/native/vision/export.py b/examples/native/vision/export.py index b10dc522ecc..047f83a10fa 100644 --- a/examples/native/vision/export.py +++ b/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,)))