From 20e83adc10fbb13d81b67ac3d57742f581af160d Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 2 Nov 2026 12:32:36 -0800 Subject: [PATCH] Match softmax input storage and layout to its output (#24437) Summary: ## Ulterior Motive Run native Swin attention softmaxes across buffer and texture storage plans. ## Rationale **What**: Repack the input of `_softmax` and `_log_softmax` when its storage or layout differs from the planned output. **Why**: The kernel reads its input using the output's storage type and texture layout. A buffer bound to a texture shader can crash; a differently packed texture can return misplaced values. ## Details - Adapt only the `self` input of these two operators. - Reuse the existing repack helper and leave matching representations unchanged. - Cover buffer-to-texture and width-to-channels-packed inputs with independent softmax references. Authored with AI using Metamate on agentcloud. Differential Revision: D129088170 --- .../runtime/vulkan/VulkanEngineExecutable.cpp | 17 ++ .../runtime/vulkan/test_vulkan_engine.cpp | 227 ++++++++++++++++++ 2 files changed, 244 insertions(+) diff --git a/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp b/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp index c50ee227eab..28657f601da 100644 --- a/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp +++ b/backends/native/runtime/vulkan/VulkanEngineExecutable.cpp @@ -1786,6 +1786,23 @@ class VulkanEngineExecutable final : public EngineExecutable { storage_at(value_id), utils::kChannelsPacked); } + // ET-VK's softmax reads its input in the output's storage type and, for + // textures, its layout. + if ((key == "aten._softmax.default" || + key == "aten._log_softmax.default") && + na.name == "self" && na.arg.kind() == ArgKind::Tensor && + valid(na.arg.as_tensor().id) && !graph_->val_is_tref(a)) { + const ValueId out_id = n.outputs.at(0).value_id; + const VkRef out = vref_at(out_id); + if (graph_->storage_type_of(a) != graph_->storage_type_of(out) || + graph_->hashed_layout_of(a) != graph_->hashed_layout_of(out)) { + a = repack( + a, + g().value(na.arg.as_tensor().id), + storage_at(out_id), + layout_at(out_id)); + } + } // 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" && diff --git a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp index 78543924005..0a3053f92ae 100644 --- a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp +++ b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp @@ -1370,6 +1370,166 @@ std::vector make_quantized_input_tensors() { return tensors; } +// result = _softmax(index(input, batch), -1) over a [2, 3, 8] input. aten.index +// keeps its operands in buffer storage, so softmax reads a buffer into a +// texture output. +std::vector make_buffered_softmax_program() { + flatbuffers::FlatBufferBuilder builder; + const auto float_meta = + create_tensor_meta(builder, fbs::ScalarType::FLOAT, {2, 3, 8}); + const std::vector> tensor_values = { + fbs::CreateTensorValueDirect(builder, "input", float_meta), + fbs::CreateTensorValueDirect( + builder, + "batch", + create_tensor_meta(builder, fbs::ScalarType::LONG, {2})), + fbs::CreateTensorValueDirect(builder, "gathered", float_meta), + fbs::CreateTensorValueDirect(builder, "result", float_meta), + }; + std::vector> nodes; + for (const char* name : {"input", "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 indices = fbs::CreateOptionalTensorListArg( + builder, + create_strings(builder, {"batch"}), + builder.CreateVector(std::vector{1})); + add_call( + "gathered", + "torch.ops.aten.index.Tensor", + { + fbs::CreateNamedArgumentDirect( + builder, "self", create_tensor_arg(builder, "input")), + fbs::CreateNamedArgumentDirect( + builder, + "indices", + fbs::CreateArgument( + builder, + fbs::ArgumentValue::OptionalTensorListArg, + indices.Union())), + }); + add_call( + "result", + "torch.ops.aten._softmax.default", + { + fbs::CreateNamedArgumentDirect( + builder, "self", create_tensor_arg(builder, "gathered")), + fbs::CreateNamedArgumentDirect( + builder, "dim", create_int_arg(builder, -1)), + fbs::CreateNamedArgumentDirect( + builder, "half_to_float", create_bool_arg(builder, false)), + }); + 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", "batch"}), + create_strings(builder, {"result"}), + builder.CreateVector(tensor_values)); + return finish_program( + builder, graph, {fbs::CreateOutputSpecDirect(builder, "result")}); +} + +// result = _softmax(linear(input, weight), -1) over a [2, 3, 8] input and a +// stored [8, 8] weight. The linear output is a width-packed texture, while the +// softmax output is planned channels-packed. +std::vector make_linear_softmax_program() { + flatbuffers::FlatBufferBuilder builder; + const auto float_meta = + create_tensor_meta(builder, fbs::ScalarType::FLOAT, {2, 3, 8}); + const auto weight_meta = + create_tensor_meta(builder, fbs::ScalarType::FLOAT, {8, 8}); + const std::vector> tensor_values = { + fbs::CreateTensorValueDirect(builder, "input", float_meta), + fbs::CreateTensorValueDirect(builder, "product", float_meta), + fbs::CreateTensorValueDirect(builder, "result", float_meta), + }; + std::vector> nodes; + for (const char* name : {"input", "weight"}) { + 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)); + }; + add_call( + "product", + "torch.ops.aten.linear.default", + { + fbs::CreateNamedArgumentDirect( + builder, "input", create_tensor_arg(builder, "input")), + fbs::CreateNamedArgumentDirect( + builder, "weight", create_tensor_arg(builder, "weight")), + fbs::CreateNamedArgumentDirect( + builder, "bias", create_none_arg(builder)), + }); + add_call( + "result", + "torch.ops.aten._softmax.default", + { + fbs::CreateNamedArgumentDirect( + builder, "self", create_tensor_arg(builder, "product")), + fbs::CreateNamedArgumentDirect( + builder, "dim", create_int_arg(builder, -1)), + fbs::CreateNamedArgumentDirect( + builder, "half_to_float", create_bool_arg(builder, false)), + }); + 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"}), + create_strings(builder, {"result"}), + builder.CreateVector(tensor_values)); + const std::vector> constants = { + fbs::CreateNamedTensorRefDirect( + builder, "weight", "weight", weight_meta, fbs::InputKind::PARAMETER), + }; + return finish_program( + builder, + graph, + {fbs::CreateOutputSpecDirect(builder, "result")}, + constants); +} + constexpr int64_t kHfTokens = 4; constexpr int64_t kHfHeads = 4; constexpr int64_t kHfKvHeads = 2; @@ -2182,6 +2342,73 @@ TEST(VulkanEngineTest, ResizesBufferViewWithItsSource) { EXPECT_EQ(output, (std::vector{4.0f, 1.0f})); } +TEST(VulkanEngineTest, ExecutesSoftmaxOnBufferPlannedInput) { + std::vector input(48); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast((i * 5) % 7) * 0.5f - 1.0f; + } + const std::vector batch = {1, 0}; + CompiledProgram compiled = compile_program(make_buffered_softmax_program()); + compiled.executable->set_input( + 0, input.data(), input.size(), ScalarType::Float); + compiled.executable->set_input( + 1, batch.data(), batch.size(), ScalarType::Long); + compiled.executable->execute(); + std::vector output(48); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::Float); + for (size_t row = 0; row < 6; ++row) { + const size_t source_row = (row + 3) % 6; + double sum = 0.0; + for (size_t i = 0; i < 8; ++i) { + sum += std::exp(input[source_row * 8 + i]); + } + for (size_t i = 0; i < 8; ++i) { + EXPECT_NEAR( + output[row * 8 + i], std::exp(input[source_row * 8 + i]) / sum, 1e-6) + << row << "," << i; + } + } +} + +TEST(VulkanEngineTest, ExecutesSoftmaxOnWidthPackedTextureInput) { + std::vector weight(64, 0.0f); + for (size_t o = 0; o < 8; ++o) { + weight[o * 8 + o] = 0.5f; + weight[o * 8 + (o + 3) % 8] = 0.25f; + } + const std::string header = + R"({"weight":{"dtype":"F32","shape":[8,8],"data_offsets":[0,256]}})"; + std::vector tensors = testing::make_safetensors(header, 256); + std::memcpy( + tensors.data() + sizeof(uint64_t) + header.size(), weight.data(), 256); + CompiledProgram compiled = + compile_program(make_linear_softmax_program(), std::move(tensors)); + std::vector input(48); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast((i * 5) % 7) * 0.5f - 1.0f; + } + compiled.executable->set_input( + 0, input.data(), input.size(), ScalarType::Float); + compiled.executable->execute(); + std::vector output(48); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::Float); + for (size_t row = 0; row < 6; ++row) { + std::vector product(8); + double sum = 0.0; + for (size_t o = 0; o < 8; ++o) { + product[o] = + 0.5 * input[row * 8 + o] + 0.25 * input[row * 8 + (o + 3) % 8]; + sum += std::exp(product[o]); + } + for (size_t o = 0; o < 8; ++o) { + EXPECT_NEAR(output[row * 8 + o], std::exp(product[o]) / sum, 1e-6) + << row << "," << o; + } + } +} + TEST(VulkanEngineTest, RejectsBufferViewResizeThatDoesNotDivide) { CompiledProgram compiled = compile_program(make_dynamic_buffer_view_index_program());