diff --git a/backends/native/runtime/vulkan/passes/InsertPrepack.cpp b/backends/native/runtime/vulkan/passes/InsertPrepack.cpp index 776cffea3a4..77f0b0b009c 100644 --- a/backends/native/runtime/vulkan/passes/InsertPrepack.cpp +++ b/backends/native/runtime/vulkan/passes/InsertPrepack.cpp @@ -19,8 +19,10 @@ namespace ptn::vulkan { namespace { bool needs_prepack(const Node& node) { - constexpr std::array kOps{ + constexpr std::array kOps{ "torch.ops.aten.add.Tensor", + "torch.ops.aten.clamp.default", + "torch.ops.aten.clone.default", "torch.ops.aten.index.Tensor", "torch.ops.aten.mul.Tensor", "torch.ops.aten.slice_copy.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..9e1f853bd63 100644 --- a/backends/native/test/runtime/vulkan/test_insert_prepack.cpp +++ b/backends/native/test/runtime/vulkan/test_insert_prepack.cpp @@ -98,6 +98,19 @@ TEST(InsertPrepackTest, PrepacksConstantSliceInput) { method.graph.node(2).input_value_ids(), std::vector({0, 4})); } +TEST(InsertPrepackTest, PrepacksConstantCloneAndClampInputs) { + for (const char* target : + {"torch.ops.aten.clone.default", "torch.ops.aten.clamp.default"}) { + SCOPED_TRACE(target); + Method method = make_method(); + method.graph.node(2).target = target; + + 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..cb2bc4a4c50 100644 --- a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp +++ b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp @@ -1370,6 +1370,87 @@ std::vector make_quantized_input_tensors() { return tensors; } +// sum = input + op(bias), where `bias` is a stored constant and op is clone, +// or with `clamp` a clamp to [-1, 50]. +std::vector make_stored_operand_program(bool clamp) { + flatbuffers::FlatBufferBuilder builder; + const auto meta = create_tensor_meta(builder, fbs::ScalarType::FLOAT, {4}); + const std::vector> tensor_values = { + fbs::CreateTensorValueDirect(builder, "input", meta), + fbs::CreateTensorValueDirect(builder, "bias", meta), + fbs::CreateTensorValueDirect(builder, "operand", meta), + fbs::CreateTensorValueDirect(builder, "sum", meta), + }; + std::vector> nodes; + for (const char* name : {"input", "bias"}) { + const std::vector> outputs = { + fbs::CreateOutputDirect(builder, name)}; + nodes.push_back(fbs::CreateNodeDirect( + builder, name, fbs::OpKind::PLACEHOLDER, "", nullptr, &outputs)); + } + const auto float_arg = [&](double value) { + return fbs::CreateArgument( + builder, + fbs::ArgumentValue::FloatArg, + fbs::CreateFloatArg(builder, value).Union()); + }; + std::vector> op_inputs = { + fbs::CreateNamedArgumentDirect( + builder, "self", create_tensor_arg(builder, "bias"))}; + if (clamp) { + op_inputs.push_back( + fbs::CreateNamedArgumentDirect(builder, "min", float_arg(-1.0))); + op_inputs.push_back( + fbs::CreateNamedArgumentDirect(builder, "max", float_arg(50.0))); + } else { + op_inputs.push_back(fbs::CreateNamedArgumentDirect( + builder, "memory_format", create_none_arg(builder))); + } + const std::vector> operand_outputs = { + fbs::CreateOutputDirect(builder, "operand")}; + nodes.push_back(fbs::CreateNodeDirect( + builder, + "operand", + fbs::OpKind::CALL_FUNCTION, + clamp ? "torch.ops.aten.clamp.default" : "torch.ops.aten.clone.default", + &op_inputs, + &operand_outputs)); + const std::vector> add_inputs = { + fbs::CreateNamedArgumentDirect( + builder, "self", create_tensor_arg(builder, "input")), + fbs::CreateNamedArgumentDirect( + builder, "other", create_tensor_arg(builder, "operand")), + fbs::CreateNamedArgumentDirect( + builder, "alpha", create_int_arg(builder, 1)), + }; + const std::vector> sum_outputs = { + fbs::CreateOutputDirect(builder, "sum")}; + nodes.push_back(fbs::CreateNodeDirect( + builder, + "sum", + fbs::OpKind::CALL_FUNCTION, + "torch.ops.aten.add.Tensor", + &add_inputs, + &sum_outputs)); + const std::vector> output_inputs = { + fbs::CreateNamedArgumentDirect( + builder, "", create_tensor_arg(builder, "sum"))}; + 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"}), + create_strings(builder, {"sum"}), + builder.CreateVector(tensor_values)); + const std::vector> constants = { + fbs::CreateNamedTensorRefDirect( + builder, "bias", "bias", meta, fbs::InputKind::PARAMETER), + }; + 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 +2263,35 @@ TEST(VulkanEngineTest, ResizesBufferViewWithItsSource) { EXPECT_EQ(output, (std::vector{4.0f, 1.0f})); } +TEST(VulkanEngineTest, ClonesAndClampsStoredConstant) { + const std::string header = + R"({"bias":{"dtype":"F32","shape":[4],"data_offsets":[0,16]}})"; + const std::vector bias = {3.0f, -2.0f, 7.0f, 100.0f}; + const std::vector input = {0.5f, 1.5f, 2.5f, 3.5f}; + for (const bool clamp : {false, true}) { + SCOPED_TRACE(clamp ? "clamp" : "clone"); + std::vector tensors = testing::make_safetensors(header, 16); + std::memcpy( + tensors.data() + sizeof(uint64_t) + header.size(), + bias.data(), + bias.size() * sizeof(float)); + CompiledProgram compiled = + compile_program(make_stored_operand_program(clamp), std::move(tensors)); + 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); + std::vector expected(4); + for (size_t i = 0; i < expected.size(); ++i) { + expected[i] = + input[i] + (clamp ? std::clamp(bias[i], -1.0f, 50.0f) : bias[i]); + } + EXPECT_EQ(output, expected); + } +} + TEST(VulkanEngineTest, RejectsBufferViewResizeThatDoesNotDivide) { CompiledProgram compiled = compile_program(make_dynamic_buffer_view_index_program());