Skip to content

Commit d974571

Browse files
Fix XNN PreLU use after free (#21480) (#21845)
Cherry-picks #21480 onto the release/1.4 branch for the 1.4.1 release. Original PR: #21480 Source commit: 0b0495d This cherry-pick PR was created with OpenAI Codex. Co-authored-by: Jake Stevens <jrstevens@meta.com>
1 parent 3e0e2dd commit d974571

4 files changed

Lines changed: 156 additions & 2 deletions

File tree

‎backends/xnnpack/runtime/XNNCompiler.cpp‎

Lines changed: 71 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include <cinttypes>
1616
#include <string>
1717
#include <unordered_map>
18+
#include <unordered_set>
1819
#include <vector>
1920

2021
#pragma clang diagnostic ignored "-Wmissing-prototypes"
@@ -1989,6 +1990,58 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
19891990
}
19901991
#undef _DEFINE
19911992

1993+
/*
1994+
Serialized id of an xvalue, or XNN_INVALID_VALUE_ID if it is not a tensor.
1995+
*/
1996+
uint32_t getSerializedValueId(ValuePtr value) noexcept {
1997+
const fb_xnnpack::XNNTensorValue* tensor_value = nullptr;
1998+
if (value->xvalue_union_type() == fb_xnnpack::XValueUnion::XNNTensorValue) {
1999+
tensor_value = value->xvalue_union_as_XNNTensorValue();
2000+
} else if (
2001+
value->xvalue_union_type() ==
2002+
fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) {
2003+
tensor_value =
2004+
value->xvalue_union_as_XNNQuantizedTensorValue()->tensor_value();
2005+
}
2006+
return tensor_value != nullptr ? tensor_value->id_out()
2007+
: XNN_INVALID_VALUE_ID;
2008+
}
2009+
2010+
/*
2011+
Serialized ids of the values that XNNPACK copies into its own packed storage
2012+
while the runtime is created. Every other constant value is referenced by raw
2013+
pointer for the lifetime of the runtime.
2014+
*/
2015+
std::unordered_set<uint32_t> getPackedValueIds(GraphPtr flatbuffer_graph) {
2016+
std::unordered_set<uint32_t> packed_ids;
2017+
auto insert_weights = [&packed_ids](uint32_t filter_id, uint32_t bias_id) {
2018+
packed_ids.insert(filter_id);
2019+
if (bias_id != XNN_INVALID_VALUE_ID) {
2020+
packed_ids.insert(bias_id);
2021+
}
2022+
};
2023+
2024+
// Deliberately an if-chain rather than a switch: XNodeUnion has ~50
2025+
// enumerators and -Wswitch-enum requires every one of them to be listed.
2026+
for (auto node : *flatbuffer_graph->xnodes()) {
2027+
auto type = node->xnode_union_type();
2028+
if (type == fb_xnnpack::XNodeUnion::XNNFullyConnected) {
2029+
auto n = node->xnode_union_as_XNNFullyConnected();
2030+
insert_weights(n->filter_id(), n->bias_id());
2031+
} else if (type == fb_xnnpack::XNodeUnion::XNNConv2d) {
2032+
auto n = node->xnode_union_as_XNNConv2d();
2033+
insert_weights(n->filter_id(), n->bias_id());
2034+
} else if (type == fb_xnnpack::XNodeUnion::XNNDepthwiseConv2d) {
2035+
auto n = node->xnode_union_as_XNNDepthwiseConv2d();
2036+
insert_weights(n->filter_id(), n->bias_id());
2037+
} else if (type == fb_xnnpack::XNodeUnion::XNNConvTranspose2d) {
2038+
auto n = node->xnode_union_as_XNNConvTranspose2d();
2039+
insert_weights(n->filter_id(), n->bias_id());
2040+
}
2041+
}
2042+
return packed_ids;
2043+
}
2044+
19922045
/*
19932046
Builds the xnnpack runtime object using the buffer pointer. The buffer pointer
19942047
must be a valid pointer to the serialized xnnpack object. It also fills the
@@ -2093,15 +2146,19 @@ ET_NODISCARD Error XNNCompiler::compileModel(
20932146
// Invalid ids do not need to be remapped
20942147
remapped_ids.emplace(XNN_INVALID_VALUE_ID, XNN_INVALID_VALUE_ID);
20952148

2096-
// If weight cache is not on we hold onto all the unpacked buffers
2097-
// and we free them at the end
2149+
// Buffers loaded from the named data map for values whose data XNNPACK packs
2150+
// during runtime creation. They stay alive until the runtime exists and are
2151+
// freed afterwards.
20982152
std::vector<FreeableBuffer> unpacked_buffers;
2153+
const std::unordered_set<uint32_t> packed_value_ids =
2154+
getPackedValueIds(flatbuffer_graph);
20992155

21002156
// External Ids for inputs and outputs
21012157
std::vector<uint32_t> input_ids;
21022158
std::vector<uint32_t> output_ids;
21032159
Error err = Error::Ok;
21042160
for (auto value : *flatbuffer_graph->xvalues()) {
2161+
size_t prev_buffers = unpacked_buffers.size();
21052162
err = defineTensor(
21062163
subgraph.get(),
21072164
remapped_ids,
@@ -2120,6 +2177,17 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21202177
if (err != Error::Ok) {
21212178
return err;
21222179
}
2180+
2181+
// Operators that don't pack (PReLU, for example) keep raw pointers into
2182+
// the constant data, so hand their buffers to the executor. A single value
2183+
// can contribute more than one buffer: a quantized weight also loads its
2184+
// scales.
2185+
if (packed_value_ids.count(getSerializedValueId(value)) == 0) {
2186+
for (size_t i = prev_buffers; i < unpacked_buffers.size(); i++) {
2187+
executor->unpacked_buffers_.push_back(std::move(unpacked_buffers[i]));
2188+
}
2189+
unpacked_buffers.resize(prev_buffers);
2190+
}
21232191
}
21242192

21252193
for (auto node : *flatbuffer_graph->xnodes()) {
@@ -2174,6 +2242,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21742242
"Failed to finalize weights cache after creating the xnn runtime");
21752243
packed_weights_names = std::move(packed_weights_names_result.get());
21762244
} else {
2245+
// XNNPACK has copied these into its own packed storage.
21772246
for (auto& buffer : unpacked_buffers) {
21782247
buffer.Free();
21792248
}

‎backends/xnnpack/runtime/XNNExecutor.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include <executorch/runtime/backend/interface.h>
1515
#include <executorch/runtime/core/error.h>
1616
#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
17+
#include <executorch/runtime/core/freeable_buffer.h>
1718

1819
#include <xnnpack.h>
1920
#include <atomic>
@@ -30,6 +31,10 @@ class XNNWeightsCache;
3031

3132
class XNNExecutor {
3233
private:
34+
// For XNN constant data that isn't packed (PreLU weights, for example),
35+
// we need to hold onto the buffers to keep them alive.
36+
std::vector<executorch::runtime::FreeableBuffer> unpacked_buffers_;
37+
3338
std::unique_ptr<xnn_runtime, decltype(&xnn_delete_runtime)> runtime_{
3439
nullptr,
3540
&xnn_delete_runtime};

‎backends/xnnpack/test/BUCK‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,13 +95,17 @@ fbcode_target(_kind = runtime.python_test,
9595
"test_xnnpack_utils.py",
9696
],
9797
deps = [
98+
"//executorch/backends/test/harness:tester",
9899
"//executorch/backends/xnnpack/partition:xnnpack_partitioner",
99100
"//executorch/backends/xnnpack/quantizer:xnnpack_quantizer",
100101
"//executorch/backends/xnnpack/test/tester:tester",
101102
"//executorch/devtools:lib",
102103
"//executorch/devtools/bundled_program:config",
103104
"//executorch/devtools/bundled_program/serialize:lib",
105+
"//executorch/exir:lib",
106+
"//executorch/exir/capture:config",
104107
"//executorch/exir/passes:constant_prop_pass",
108+
"//executorch/runtime:runtime",
105109
"//pytorch/ao:torchao", # @manual
106110
],
107111
external_deps = [

‎backends/xnnpack/test/ops/test_prelu.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,18 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7+
import gc
8+
import tempfile
79
import unittest
10+
from pathlib import Path
811

912
import torch
13+
from executorch.backends.test.harness.stages import StageType
14+
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
1015
from executorch.backends.xnnpack.test.tester import Tester
16+
from executorch.exir import to_edge_transform_and_lower
17+
from executorch.exir.capture._config import ExecutorchBackendConfig
18+
from executorch.runtime import Runtime, Verification
1119

1220

1321
class TestPrelu(unittest.TestCase):
@@ -23,6 +31,14 @@ def forward(self, x):
2331
a = self.prelu(x)
2432
return a
2533

34+
class ConstWPrelu(torch.nn.Module):
35+
def __init__(self):
36+
super().__init__()
37+
self.register_buffer("w", torch.ones(3, dtype=torch.float32))
38+
39+
def forward(self, x):
40+
return torch.ops.aten.prelu.default(x, self.w)
41+
2642
def _test_prelu(self, module, inputs):
2743
(
2844
Tester(module, inputs)
@@ -38,6 +54,26 @@ def _test_prelu(self, module, inputs):
3854
.run_method_and_compare_outputs()
3955
)
4056

57+
def _load_and_compare_from_file(self, write_program, inputs, expected):
58+
with tempfile.TemporaryDirectory() as temp_dir:
59+
pte_path = Path(temp_dir) / "prelu.pte"
60+
with pte_path.open("wb") as f:
61+
write_program(f)
62+
63+
rt = Runtime.get()
64+
program = rt.load_program(pte_path, verification=Verification.Minimal)
65+
method = program.load_method("forward")
66+
actual = method.execute(inputs)[0].clone()
67+
68+
# The program mmaps the .pte and keeps it mapped for as long as it
69+
# or any of its methods is alive. Windows refuses to delete a
70+
# mapped file, so release everything before the temp dir is
71+
# cleaned up.
72+
del method, program
73+
gc.collect()
74+
75+
self.assertTrue(torch.allclose(expected, actual, atol=1e-5))
76+
4177
@unittest.skip("XNNPACK Expects FP16 inputs but FP32 weights")
4278
def _test_fp16_prelu(self):
4379
module = self.PReLU().to(torch.float16)
@@ -48,3 +84,43 @@ def test_fp32_prelu(self):
4884
module = self.PReLU()
4985
inputs = (torch.randn(1, 5, 3, 2),)
5086
self._test_prelu(module, inputs)
87+
88+
def test_fp32_prelu_file_load(self):
89+
"""
90+
Make sure that PreLU doesn't free its weight buffer after load. It's a weird
91+
op that doesn't copy or pack its data, so we need to hold onto the buffer.
92+
Run specifically from a file to exercise the path.
93+
"""
94+
module = self.PReLU()
95+
module.eval()
96+
x = torch.randn(1, 5, 3, 2)
97+
expected = module(x)
98+
99+
tester = Tester(module, (x,))
100+
tester.export()
101+
tester.to_edge_transform_and_lower()
102+
tester.check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
103+
tester.to_executorch()
104+
tester.serialize()
105+
106+
buf = tester.stages[StageType.SERIALIZE].artifact
107+
self._load_and_compare_from_file(lambda f: f.write(buf), (x,), expected)
108+
109+
def test_fp32_prelu_constant_weight_empty_decompositions_file_load(self):
110+
module = self.ConstWPrelu().eval()
111+
x = torch.randn(2, 3, 3, 3, device="cpu", dtype=torch.float32)
112+
expected = module(x)
113+
114+
exported = torch.export.export(module, args=(x,), strict=True)
115+
exported = exported.run_decompositions({})
116+
117+
edge_pm = to_edge_transform_and_lower(
118+
exported,
119+
partitioner=[XnnpackPartitioner()],
120+
compile_config=None,
121+
)
122+
et_pm = edge_pm.to_executorch(
123+
config=ExecutorchBackendConfig(extract_delegate_segments=True)
124+
)
125+
126+
self._load_and_compare_from_file(et_pm.write_to_file, (x,), expected)

0 commit comments

Comments
 (0)