Fix XNN PreLU use after free (#21480)

Summary:

Most XNNPACK parameter / constant data is packed and XNNPACK takes ownership of it. This means that we can free the buffer after load. However, PreLU doesn't pack and thus expects the buffer to stay alive. There is a reproducible segfault when running PreLU on ET XNNPACK when running from a file instead of in-memory buffer data.

This change fixes that by only freeing buffers for ops we know take ownership of the data. Practically speaking this should be a no-op for most cases like linear and conv, but it plays it safe for other constant tensors.

Fixes #17559.

Differential Revision: D114141346

Author: JakeStevens

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"
@@ -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<FreeableBuffer> 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<uint32_t, size_t> named_data_buffer_map;
+
   // 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()) {
+  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<uint32_t> 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]));
+        }
+      }
     }
   }
 

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

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