From 25bf43e3342fc332233a01aeaa4a58b6cc9f2903 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 2 Nov 2026 12:32:44 -0800 Subject: [PATCH] Copy slice-scatter inputs into channels-packed textures (#24438) Summary: ## Ulterior Motive Unblock Swin V2's runtime query/key/value bias update. ## Rationale **What**: Copy and repack `slice_scatter` inputs and output in the kernel's required channels-packed texture representation. **Why**: The stored bias is cloned and updated at runtime, but its planned buffers do not satisfy the operator's texture-only contract. ## Details - Repack `self` and `src` only when they are not standard channels-packed textures. - Reject over-limit tensors before constructing temporary tensors. - Dispatch into a matching temporary and copy back when the planned output differs. Authored with AI using Metamate on agentcloud. Differential Revision: D129088172 --- .../runtime/vulkan/VulkanEngineExecutable.cpp | 40 ++++++ .../runtime/vulkan/test_vulkan_engine.cpp | 123 ++++++++++++++++++ 2 files changed, 163 insertions(+) diff --git a/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp b/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp index c50ee227eab..df440f6f425 100644 --- a/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp +++ b/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp @@ -1749,6 +1749,24 @@ class VulkanEngineExecutable final : public EngineExecutable { } } + // 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 ids = n.input_value_ids(); + for_each_output_value(n, [&](ValueId id) { ids.push_back(id); }); + if (std::ranges::any_of(ids, [&](ValueId id) { + return 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 args; args.reserve(n.inputs.size() + 1); for (const NamedArgument& na : n.inputs) { @@ -1786,6 +1804,16 @@ class VulkanEngineExecutable final : public EngineExecutable { 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 +1856,18 @@ class VulkanEngineExecutable final : public EngineExecutable { std::vector 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") && diff --git a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp index 78543924005..7a961696a46 100644 --- a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp +++ b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp @@ -1370,6 +1370,103 @@ std::vector make_quantized_input_tensors() { 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 make_buffered_slice_scatter_program() { + flatbuffers::FlatBufferBuilder builder; + const auto meta = create_tensor_meta(builder, fbs::ScalarType::FLOAT, {2, 6}); + const std::vector> 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> nodes; + for (const char* name : {"input", "src", "batch"}) { + const std::vector> 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, + const std::vector>& inputs) { + const std::vector> 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{1})); + return std::vector>{ + 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> 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 +2279,32 @@ TEST(VulkanEngineTest, ResizesBufferViewWithItsSource) { EXPECT_EQ(output, (std::vector{4.0f, 1.0f})); } +TEST(VulkanEngineTest, ExecutesSliceScatterOnBufferPlannedTensors) { + std::vector input(12); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(i + 1); + } + const std::vector src = {-1.0f, -2.0f, -3.0f, -4.0f}; + const std::vector 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 output(12); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::Float); + std::vector 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());