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