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"
@@ -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