Fix XNN PreLU use after free (#21480)

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -15,6 +15,7 @@
 #include <cinttypes>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #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<uint32_t> getPackedValueIds(GraphPtr flatbuffer_graph) {
+  std::unordered_set<uint32_t> 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<FreeableBuffer> unpacked_buffers;
+  const std::unordered_set<uint32_t> packed_value_ids =
+      getPackedValueIds(flatbuffer_graph);
 
   // External Ids for inputs and outputs
   std::vector<uint32_t> input_ids;
   std::vector<uint32_t> 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();
     }

backends/xnnpack/runtime/XNNExecutor.h

@@ -14,6 +14,7 @@
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/core/freeable_buffer.h>
 
 #include <xnnpack.h>
 #include <atomic>
@@ -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<executorch::runtime::FreeableBuffer> unpacked_buffers_;
+
   std::unique_ptr<xnn_runtime, decltype(&xnn_delete_runtime)> runtime_{
       nullptr,
       &xnn_delete_runtime};

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 = [

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)