From d5dcc83f3c99626f79b188eba0bf19c6503379c5 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Thu, 30 Jul 2026 05:53:03 -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. Differential Revision: D114141346 --- backends/xnnpack/runtime/XNNCompiler.cpp | 69 +++++++++++++++++++++--- backends/xnnpack/runtime/XNNExecutor.h | 5 ++ backends/xnnpack/test/ops/test_prelu.py | 67 +++++++++++++++++++++++ backends/xnnpack/xnnpack_preprocess.py | 1 + 4 files changed, 136 insertions(+), 6 deletions(-) diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 3f2fb5a1aaa..ce38035521d 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" @@ -2093,19 +2094,27 @@ 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. After xnn_create_runtime, + // buffers consumed by packing operators are freed; the rest are moved + // into the executor to keep them alive for non-packing operators. std::vector unpacked_buffers; + // Maps xvalue index to unpacked_buffers index for values whose data was + // loaded from the named data map. Used to selectively retain buffers that + // are still referenced at runtime (non-packing operators). + std::unordered_map named_data_buffer_map; + // External Ids for inputs and outputs std::vector input_ids; std::vector output_ids; Error err = Error::Ok; - for (auto value : *flatbuffer_graph->xvalues()) { + auto xvalues = flatbuffer_graph->xvalues(); + for (uint32_t i = 0; i < xvalues->size(); i++) { + size_t prev_buffers = unpacked_buffers.size(); err = defineTensor( subgraph.get(), remapped_ids, - value, + xvalues->Get(i), flatbuffer_graph, constant_data, constant_data_size, @@ -2120,6 +2129,10 @@ ET_NODISCARD Error XNNCompiler::compileModel( if (err != Error::Ok) { return err; } + + if (unpacked_buffers.size() > prev_buffers) { + named_data_buffer_map[i] = prev_buffers; + } } for (auto node : *flatbuffer_graph->xnodes()) { @@ -2174,8 +2187,52 @@ 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 { - for (auto& buffer : unpacked_buffers) { - buffer.Free(); + // Operators like convolution and fully-connected pack weights during load, + // so those buffers can be freed. Other operators (PreLU) retain raw + // pointers to the original constant data, so those buffers need to remain + // alive. + if (!named_data_buffer_map.empty()) { + std::unordered_set packed_value_indices; + for (auto node : *flatbuffer_graph->xnodes()) { + auto type = node->xnode_union_type(); + switch (type) { + case fb_xnnpack::XNodeUnion::XNNFullyConnected: { + auto n = node->xnode_union_as_XNNFullyConnected(); + packed_value_indices.insert(n->filter_id()); + packed_value_indices.insert(n->bias_id()); + break; + } + case fb_xnnpack::XNodeUnion::XNNConv2d: { + auto n = node->xnode_union_as_XNNConv2d(); + packed_value_indices.insert(n->filter_id()); + packed_value_indices.insert(n->bias_id()); + break; + } + case fb_xnnpack::XNodeUnion::XNNDepthwiseConv2d: { + auto n = node->xnode_union_as_XNNDepthwiseConv2d(); + packed_value_indices.insert(n->filter_id()); + packed_value_indices.insert(n->bias_id()); + break; + } + case fb_xnnpack::XNodeUnion::XNNConvTranspose2d: { + auto n = node->xnode_union_as_XNNConvTranspose2d(); + packed_value_indices.insert(n->filter_id()); + packed_value_indices.insert(n->bias_id()); + break; + } + default: + break; + } + } + + for (auto& [value_idx, buffer_idx] : named_data_buffer_map) { + if (packed_value_indices.count(value_idx)) { + unpacked_buffers[buffer_idx].Free(); + } else { + executor->unpacked_buffers_.push_back( + std::move(unpacked_buffers[buffer_idx])); + } + } } } 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/ops/test_prelu.py b/backends/xnnpack/test/ops/test_prelu.py index 47b2851278c..af49df5b726 100644 --- a/backends/xnnpack/test/ops/test_prelu.py +++ b/backends/xnnpack/test/ops/test_prelu.py @@ -4,10 +4,17 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +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 +30,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 +53,18 @@ 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] + 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 +75,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) diff --git a/backends/xnnpack/xnnpack_preprocess.py b/backends/xnnpack/xnnpack_preprocess.py index cdceb8a90a1..2cf6b88cc46 100644 --- a/backends/xnnpack/xnnpack_preprocess.py +++ b/backends/xnnpack/xnnpack_preprocess.py @@ -140,6 +140,7 @@ def preprocess( passes.append(ConvertToLinearPass) passes = passes if len(passes) > 0 else None + # XNNPACK Delegate Specific Passes ep = XNNPACKPassManager(ep, passes=passes).transform() graph_module = ep.graph_module