Add Python runtime end-to-end tests for XNNPACK on Linux

runtime/test/TARGETS

@@ -21,3 +21,14 @@ runtime.python_test(
         "//executorch/devtools/etdump:serialize",
     ],
 )
+
+runtime.python_test(
+    name = "test_runtime_xnnpack",
+    srcs = ["test_runtime_xnnpack.py"],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
+        "//executorch/exir:lib",
+        "//executorch/runtime:runtime",
+    ],
+)

runtime/test/test_runtime_xnnpack.py

@@ -0,0 +1,181 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Python runtime end-to-end tests for the XNNPACK backend on Linux.
+
+Covers the export -> .pte -> Python runtime flow that developers use to
+validate exported models before deploying to device. Placing these tests
+under ``runtime/test`` lets the standard unittest jobs collect them cleanly.
+
+See https://github.com/pytorch/executorch/issues/11225
+"""
+
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+from executorch.runtime import Runtime, Verification
+from torch.export import export
+
+
+def _export_and_execute(
+    model: torch.nn.Module, example_inputs: tuple[torch.Tensor, ...]
+):
+    """Export *model* with XNNPACK, save to a temp .pte, and run it via Runtime."""
+    with tempfile.TemporaryDirectory() as temp_dir, torch.no_grad():
+        model.eval()
+        aten = export(model, example_inputs, strict=True)
+        edge = to_edge_transform_and_lower(
+            aten,
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+            partitioner=[XnnpackPartitioner()],
+        )
+        et = edge.to_executorch()
+
+        pte_path = Path(temp_dir) / "xnnpack_runtime_test.pte"
+        et.save(str(pte_path))
+
+        runtime = Runtime.get()
+        program = runtime.load_program(pte_path, verification=Verification.Minimal)
+        method = program.load_method("forward")
+        assert method is not None, "forward method should exist in exported program"
+        return method.execute(example_inputs)
+
+
+@unittest.skipUnless(
+    sys.platform == "linux",
+    "XNNPACK Python runtime end-to-end coverage in this batch targets Linux only",
+)
+class RuntimeXNNPACKTest(unittest.TestCase):
+    """Export → .pte → Python Runtime tests for XNNPACK on Linux."""
+
+    # ------------------------------------------------------------------
+    # Simple arithmetic
+    # ------------------------------------------------------------------
+    def test_add(self):
+        class Add(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        model = Add()
+        inputs = (torch.randn(2, 3), torch.randn(2, 3))
+
+        expected = model(*inputs)
+        actual = _export_and_execute(model, inputs)
+        torch.testing.assert_close(actual[0], expected, atol=1e-4, rtol=1e-4)
+
+    # ------------------------------------------------------------------
+    # Linear layer (fp32)
+    # ------------------------------------------------------------------
+    def test_linear(self):
+        model = torch.nn.Linear(16, 8)
+        inputs = (torch.randn(4, 16),)
+
+        with torch.no_grad():
+            expected = model(*inputs)
+        actual = _export_and_execute(model, inputs)
+        torch.testing.assert_close(actual[0], expected, atol=1e-4, rtol=1e-4)
+
+    # ------------------------------------------------------------------
+    # Conv2d + ReLU (common vision pattern)
+    # ------------------------------------------------------------------
+    def test_conv2d_relu(self):
+        model = torch.nn.Sequential(
+            torch.nn.Conv2d(3, 16, 3, padding=1),
+            torch.nn.ReLU(),
+        )
+        inputs = (torch.randn(1, 3, 8, 8),)
+
+        with torch.no_grad():
+            expected = model(*inputs)
+        actual = _export_and_execute(model, inputs)
+        torch.testing.assert_close(actual[0], expected, atol=1e-3, rtol=1e-3)
+
+    # ------------------------------------------------------------------
+    # Small MLP (multiple linear + activation)
+    # ------------------------------------------------------------------
+    def test_mlp(self):
+        model = torch.nn.Sequential(
+            torch.nn.Linear(32, 64),
+            torch.nn.ReLU(),
+            torch.nn.Linear(64, 10),
+        )
+        inputs = (torch.randn(2, 32),)
+
+        with torch.no_grad():
+            expected = model(*inputs)
+        actual = _export_and_execute(model, inputs)
+        torch.testing.assert_close(actual[0], expected, atol=1e-3, rtol=1e-3)
+
+    # ------------------------------------------------------------------
+    # BatchNorm + Conv (common in MobileNet-style models)
+    # Skipped: FuseBatchNormPass crashes on Sequential(Conv2d, BN) export.
+    # TODO(#11225): re-enable once the XNNPACK pass is fixed.
+    # ------------------------------------------------------------------
+    @unittest.skip("FuseBatchNormPass bug in XNNPACK backend")
+    def test_conv_bn(self):
+        model = torch.nn.Sequential(
+            torch.nn.Conv2d(3, 16, 3, padding=1),
+            torch.nn.BatchNorm2d(16),
+            torch.nn.ReLU(),
+        )
+        model.eval()
+        inputs = (torch.randn(1, 3, 8, 8),)
+
+        with torch.no_grad():
+            expected = model(*inputs)
+        actual = _export_and_execute(model, inputs)
+        torch.testing.assert_close(actual[0], expected, atol=1e-3, rtol=1e-3)
+
+    # ------------------------------------------------------------------
+    # Depthwise separable conv (MobileNet building block)
+    # ------------------------------------------------------------------
+    def test_depthwise_separable_conv(self):
+        model = torch.nn.Sequential(
+            # Depthwise
+            torch.nn.Conv2d(16, 16, 3, padding=1, groups=16),
+            torch.nn.ReLU(),
+            # Pointwise
+            torch.nn.Conv2d(16, 32, 1),
+            torch.nn.ReLU(),
+        )
+        inputs = (torch.randn(1, 16, 8, 8),)
+
+        with torch.no_grad():
+            expected = model(*inputs)
+        actual = _export_and_execute(model, inputs)
+        torch.testing.assert_close(actual[0], expected, atol=1e-3, rtol=1e-3)
+
+    # ------------------------------------------------------------------
+    # Avgpool + Flatten + Linear (classifier head)
+    # ------------------------------------------------------------------
+    def test_classifier_head(self):
+        class ClassifierHead(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.pool = torch.nn.AdaptiveAvgPool2d(1)
+                self.fc = torch.nn.Linear(32, 10)
+
+            def forward(self, x):
+                x = self.pool(x)
+                x = x.flatten(1)
+                return self.fc(x)
+
+        model = ClassifierHead()
+        inputs = (torch.randn(1, 32, 8, 8),)
+
+        with torch.no_grad():
+            expected = model(*inputs)
+        actual = _export_and_execute(model, inputs)
+        torch.testing.assert_close(actual[0], expected, atol=1e-3, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()