Skip to content

Commit ad3a71f

Browse files
authored
Fix XNN PreLU use after free (pytorch#21480)
Differential Revision: D114141346 Pull Request resolved: pytorch#21480
1 parent f67adf5 commit ad3a71f

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"
@@ -2064,6 +2065,58 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
20642065
}
20652066
#undef _DEFINE
20662067

2068+
/*
2069+
Serialized id of an xvalue, or XNN_INVALID_VALUE_ID if it is not a tensor.
2070+
*/
2071+
uint32_t getSerializedValueId(ValuePtr value) noexcept {
2072+
const fb_xnnpack::XNNTensorValue* tensor_value = nullptr;
2073+
if (value->xvalue_union_type() == fb_xnnpack::XValueUnion::XNNTensorValue) {
2074+
tensor_value = value->xvalue_union_as_XNNTensorValue();
2075+
} else if (
2076+
value->xvalue_union_type() ==
2077+
fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) {
2078+
tensor_value =
2079+
value->xvalue_union_as_XNNQuantizedTensorValue()->tensor_value();
2080+
}
2081+
return tensor_value != nullptr ? tensor_value->id_out()
2082+
: XNN_INVALID_VALUE_ID;
2083+
}
2084+
2085+
/*
2086+
Serialized ids of the values that XNNPACK copies into its own packed storage
2087+
while the runtime is created. Every other constant value is referenced by raw
2088+
pointer for the lifetime of the runtime.
2089+
*/
2090+
std::unordered_set<uint32_t> getPackedValueIds(GraphPtr flatbuffer_graph) {
2091+
std::unordered_set<uint32_t> packed_ids;
2092+
auto insert_weights = [&packed_ids](uint32_t filter_id, uint32_t bias_id) {
2093+
packed_ids.insert(filter_id);
2094+
if (bias_id != XNN_INVALID_VALUE_ID) {
2095+
packed_ids.insert(bias_id);
2096+
}
2097+
};
2098+
2099+
// Deliberately an if-chain rather than a switch: XNodeUnion has ~50
2100+
// enumerators and -Wswitch-enum requires every one of them to be listed.
2101+
for (auto node : *flatbuffer_graph->xnodes()) {
2102+
auto type = node->xnode_union_type();
2103+
if (type == fb_xnnpack::XNodeUnion::XNNFullyConnected) {
2104+
auto n = node->xnode_union_as_XNNFullyConnected();
2105+
insert_weights(n->filter_id(), n->bias_id());
2106+
} else if (type == fb_xnnpack::XNodeUnion::XNNConv2d) {
2107+
auto n = node->xnode_union_as_XNNConv2d();
2108+
insert_weights(n->filter_id(), n->bias_id());
2109+
} else if (type == fb_xnnpack::XNodeUnion::XNNDepthwiseConv2d) {
2110+
auto n = node->xnode_union_as_XNNDepthwiseConv2d();
2111+
insert_weights(n->filter_id(), n->bias_id());
2112+
} else if (type == fb_xnnpack::XNodeUnion::XNNConvTranspose2d) {
2113+
auto n = node->xnode_union_as_XNNConvTranspose2d();
2114+
insert_weights(n->filter_id(), n->bias_id());
2115+
}
2116+
}
2117+
return packed_ids;
2118+
}
2119+
20672120
/*
20682121
Builds the xnnpack runtime object using the buffer pointer. The buffer pointer
20692122
must be a valid pointer to the serialized xnnpack object. It also fills the
@@ -2168,15 +2221,19 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21682221
// Invalid ids do not need to be remapped
21692222
remapped_ids.emplace(XNN_INVALID_VALUE_ID, XNN_INVALID_VALUE_ID);
21702223

2171-
// If weight cache is not on we hold onto all the unpacked buffers
2172-
// and we free them at the end
2224+
// Buffers loaded from the named data map for values whose data XNNPACK packs
2225+
// during runtime creation. They stay alive until the runtime exists and are
2226+
// freed afterwards.
21732227
std::vector<FreeableBuffer> unpacked_buffers;
2228+
const std::unordered_set<uint32_t> packed_value_ids =
2229+
getPackedValueIds(flatbuffer_graph);
21742230

21752231
// External Ids for inputs and outputs
21762232
std::vector<uint32_t> input_ids;
21772233
std::vector<uint32_t> output_ids;
21782234
Error err = Error::Ok;
21792235
for (auto value : *flatbuffer_graph->xvalues()) {
2236+
size_t prev_buffers = unpacked_buffers.size();
21802237
err = defineTensor(
21812238
subgraph.get(),
21822239
remapped_ids,
@@ -2195,6 +2252,17 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21952252
if (err != Error::Ok) {
21962253
return err;
21972254
}
2255+
2256+
// Operators that don't pack (PReLU, for example) keep raw pointers into
2257+
// the constant data, so hand their buffers to the executor. A single value
2258+
// can contribute more than one buffer: a quantized weight also loads its
2259+
// scales.
2260+
if (packed_value_ids.count(getSerializedValueId(value)) == 0) {
2261+
for (size_t i = prev_buffers; i < unpacked_buffers.size(); i++) {
2262+
executor->unpacked_buffers_.push_back(std::move(unpacked_buffers[i]));
2263+
}
2264+
unpacked_buffers.resize(prev_buffers);
2265+
}
21982266
}
21992267

22002268
for (auto node : *flatbuffer_graph->xnodes()) {
@@ -2249,6 +2317,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
22492317
"Failed to finalize weights cache after creating the xnn runtime");
22502318
packed_weights_names = std::move(packed_weights_names_result.get());
22512319
} else {
2320+
// XNNPACK has copied these into its own packed storage.
22522321
for (auto& buffer : unpacked_buffers) {
22532322
buffer.Free();
22542323
}

‎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)