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)