From 0b0495dff88a00204d26c445bb60b30d60044533 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 4 Aug 2026 07:49:51 -0700 Subject: [PATCH] Fix XNN PreLU use after free (#21480) Summary: Most XNNPACK parameter / constant data is packed and XNNPACK takes ownership of it. This means that we can free the buffer after load. However, PreLU doesn't pack and thus expects the buffer to stay alive. There is a reproducible segfault when running PreLU on ET XNNPACK when running from a file instead of in-memory buffer data. This change fixes that by only freeing buffers for ops we know take ownership of the data. Practically speaking this should be a no-op for most cases like linear and conv, but it plays it safe for other constant tensors. Fixes https://github.com/pytorch/executorch/issues/17559. Reviewed By: digantdesai Differential Revision: D114141346 --- backends/xnnpack/runtime/XNNCompiler.cpp | 73 ++++++++++++++++++++++- backends/xnnpack/runtime/XNNExecutor.h | 5 ++ backends/xnnpack/test/BUCK | 4 ++ backends/xnnpack/test/ops/test_prelu.py | 76 ++++++++++++++++++++++++ 4 files changed, 156 insertions(+), 2 deletions(-) diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 96f3316999c..2b8ab9b3181 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #pragma clang diagnostic ignored "-Wmissing-prototypes" @@ -2064,6 +2065,58 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) { } #undef _DEFINE +/* +Serialized id of an xvalue, or XNN_INVALID_VALUE_ID if it is not a tensor. +*/ +uint32_t getSerializedValueId(ValuePtr value) noexcept { + const fb_xnnpack::XNNTensorValue* tensor_value = nullptr; + if (value->xvalue_union_type() == fb_xnnpack::XValueUnion::XNNTensorValue) { + tensor_value = value->xvalue_union_as_XNNTensorValue(); + } else if ( + value->xvalue_union_type() == + fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) { + tensor_value = + value->xvalue_union_as_XNNQuantizedTensorValue()->tensor_value(); + } + return tensor_value != nullptr ? tensor_value->id_out() + : XNN_INVALID_VALUE_ID; +} + +/* +Serialized ids of the values that XNNPACK copies into its own packed storage +while the runtime is created. Every other constant value is referenced by raw +pointer for the lifetime of the runtime. +*/ +std::unordered_set getPackedValueIds(GraphPtr flatbuffer_graph) { + std::unordered_set packed_ids; + auto insert_weights = [&packed_ids](uint32_t filter_id, uint32_t bias_id) { + packed_ids.insert(filter_id); + if (bias_id != XNN_INVALID_VALUE_ID) { + packed_ids.insert(bias_id); + } + }; + + // Deliberately an if-chain rather than a switch: XNodeUnion has ~50 + // enumerators and -Wswitch-enum requires every one of them to be listed. + for (auto node : *flatbuffer_graph->xnodes()) { + auto type = node->xnode_union_type(); + if (type == fb_xnnpack::XNodeUnion::XNNFullyConnected) { + auto n = node->xnode_union_as_XNNFullyConnected(); + insert_weights(n->filter_id(), n->bias_id()); + } else if (type == fb_xnnpack::XNodeUnion::XNNConv2d) { + auto n = node->xnode_union_as_XNNConv2d(); + insert_weights(n->filter_id(), n->bias_id()); + } else if (type == fb_xnnpack::XNodeUnion::XNNDepthwiseConv2d) { + auto n = node->xnode_union_as_XNNDepthwiseConv2d(); + insert_weights(n->filter_id(), n->bias_id()); + } else if (type == fb_xnnpack::XNodeUnion::XNNConvTranspose2d) { + auto n = node->xnode_union_as_XNNConvTranspose2d(); + insert_weights(n->filter_id(), n->bias_id()); + } + } + return packed_ids; +} + /* Builds the xnnpack runtime object using the buffer pointer. The buffer pointer must be a valid pointer to the serialized xnnpack object. It also fills the @@ -2168,15 +2221,19 @@ ET_NODISCARD Error XNNCompiler::compileModel( // Invalid ids do not need to be remapped remapped_ids.emplace(XNN_INVALID_VALUE_ID, XNN_INVALID_VALUE_ID); - // If weight cache is not on we hold onto all the unpacked buffers - // and we free them at the end + // Buffers loaded from the named data map for values whose data XNNPACK packs + // during runtime creation. They stay alive until the runtime exists and are + // freed afterwards. std::vector unpacked_buffers; + const std::unordered_set packed_value_ids = + getPackedValueIds(flatbuffer_graph); // External Ids for inputs and outputs std::vector input_ids; std::vector output_ids; Error err = Error::Ok; for (auto value : *flatbuffer_graph->xvalues()) { + size_t prev_buffers = unpacked_buffers.size(); err = defineTensor( subgraph.get(), remapped_ids, @@ -2195,6 +2252,17 @@ ET_NODISCARD Error XNNCompiler::compileModel( if (err != Error::Ok) { return err; } + + // Operators that don't pack (PReLU, for example) keep raw pointers into + // the constant data, so hand their buffers to the executor. A single value + // can contribute more than one buffer: a quantized weight also loads its + // scales. + if (packed_value_ids.count(getSerializedValueId(value)) == 0) { + for (size_t i = prev_buffers; i < unpacked_buffers.size(); i++) { + executor->unpacked_buffers_.push_back(std::move(unpacked_buffers[i])); + } + unpacked_buffers.resize(prev_buffers); + } } for (auto node : *flatbuffer_graph->xnodes()) { @@ -2249,6 +2317,7 @@ ET_NODISCARD Error XNNCompiler::compileModel( "Failed to finalize weights cache after creating the xnn runtime"); packed_weights_names = std::move(packed_weights_names_result.get()); } else { + // XNNPACK has copied these into its own packed storage. for (auto& buffer : unpacked_buffers) { buffer.Free(); } diff --git a/backends/xnnpack/runtime/XNNExecutor.h b/backends/xnnpack/runtime/XNNExecutor.h index fd48c47a9cf..7d966b0f6a5 100644 --- a/backends/xnnpack/runtime/XNNExecutor.h +++ b/backends/xnnpack/runtime/XNNExecutor.h @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -30,6 +31,10 @@ class XNNWeightsCache; class XNNExecutor { private: + // For XNN constant data that isn't packed (PreLU weights, for example), + // we need to hold onto the buffers to keep them alive. + std::vector unpacked_buffers_; + std::unique_ptr runtime_{ nullptr, &xnn_delete_runtime}; diff --git a/backends/xnnpack/test/BUCK b/backends/xnnpack/test/BUCK index 3ffe6fbf53d..943a4f87c44 100644 --- a/backends/xnnpack/test/BUCK +++ b/backends/xnnpack/test/BUCK @@ -95,13 +95,17 @@ fbcode_target(_kind = runtime.python_test, "test_xnnpack_utils.py", ], deps = [ + "//executorch/backends/test/harness:tester", "//executorch/backends/xnnpack/partition:xnnpack_partitioner", "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer", "//executorch/backends/xnnpack/test/tester:tester", "//executorch/devtools:lib", "//executorch/devtools/bundled_program:config", "//executorch/devtools/bundled_program/serialize:lib", + "//executorch/exir:lib", + "//executorch/exir/capture:config", "//executorch/exir/passes:constant_prop_pass", + "//executorch/runtime:runtime", "//pytorch/ao:torchao", # @manual ], external_deps = [ diff --git a/backends/xnnpack/test/ops/test_prelu.py b/backends/xnnpack/test/ops/test_prelu.py index 47b2851278c..86441b399e9 100644 --- a/backends/xnnpack/test/ops/test_prelu.py +++ b/backends/xnnpack/test/ops/test_prelu.py @@ -4,10 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import gc +import tempfile import unittest +from pathlib import Path import torch +from executorch.backends.test.harness.stages import StageType +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner from executorch.backends.xnnpack.test.tester import Tester +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.capture._config import ExecutorchBackendConfig +from executorch.runtime import Runtime, Verification class TestPrelu(unittest.TestCase): @@ -23,6 +31,14 @@ def forward(self, x): a = self.prelu(x) return a + class ConstWPrelu(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.ones(3, dtype=torch.float32)) + + def forward(self, x): + return torch.ops.aten.prelu.default(x, self.w) + def _test_prelu(self, module, inputs): ( Tester(module, inputs) @@ -38,6 +54,26 @@ def _test_prelu(self, module, inputs): .run_method_and_compare_outputs() ) + def _load_and_compare_from_file(self, write_program, inputs, expected): + with tempfile.TemporaryDirectory() as temp_dir: + pte_path = Path(temp_dir) / "prelu.pte" + with pte_path.open("wb") as f: + write_program(f) + + rt = Runtime.get() + program = rt.load_program(pte_path, verification=Verification.Minimal) + method = program.load_method("forward") + actual = method.execute(inputs)[0].clone() + + # The program mmaps the .pte and keeps it mapped for as long as it + # or any of its methods is alive. Windows refuses to delete a + # mapped file, so release everything before the temp dir is + # cleaned up. + del method, program + gc.collect() + + self.assertTrue(torch.allclose(expected, actual, atol=1e-5)) + @unittest.skip("XNNPACK Expects FP16 inputs but FP32 weights") def _test_fp16_prelu(self): module = self.PReLU().to(torch.float16) @@ -48,3 +84,43 @@ def test_fp32_prelu(self): module = self.PReLU() inputs = (torch.randn(1, 5, 3, 2),) self._test_prelu(module, inputs) + + def test_fp32_prelu_file_load(self): + """ + Make sure that PreLU doesn't free its weight buffer after load. It's a weird + op that doesn't copy or pack its data, so we need to hold onto the buffer. + Run specifically from a file to exercise the path. + """ + module = self.PReLU() + module.eval() + x = torch.randn(1, 5, 3, 2) + expected = module(x) + + tester = Tester(module, (x,)) + tester.export() + tester.to_edge_transform_and_lower() + tester.check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + tester.to_executorch() + tester.serialize() + + buf = tester.stages[StageType.SERIALIZE].artifact + self._load_and_compare_from_file(lambda f: f.write(buf), (x,), expected) + + def test_fp32_prelu_constant_weight_empty_decompositions_file_load(self): + module = self.ConstWPrelu().eval() + x = torch.randn(2, 3, 3, 3, device="cpu", dtype=torch.float32) + expected = module(x) + + exported = torch.export.export(module, args=(x,), strict=True) + exported = exported.run_decompositions({}) + + edge_pm = to_edge_transform_and_lower( + exported, + partitioner=[XnnpackPartitioner()], + compile_config=None, + ) + et_pm = edge_pm.to_executorch( + config=ExecutorchBackendConfig(extract_delegate_segments=True) + ) + + self._load_and_compare_from_file(et_pm.write_to_file, (x,), expected)