introduce model-level end2end tests to dim order tests with different delegate (#6093)

Gasoonjia · facebook-github-bot · commit 15d872e93481 · 2024-10-10T10:31:51.000-07:00
Summary:

This diff introduced end2end tests on several models + delegation combinations.

Models: llama2, resnet18, mobilenet_v3
Delegate: no delegate, xnnpack

Differential Revision: D64174329
diff --git a/exir/tests/TARGETS b/exir/tests/TARGETS
@@ -376,7 +376,9 @@ python_unittest(
     deps = [
         ":test_memory_format_ops_pass_utils",
         "//caffe2:torch",
+        "//executorch/examples/models/llama2:llama2_model",
         "//executorch/extension/pybindings:aten_lib",  # @manual
+        "//pytorch/vision:torchvision",  # @manual
     ],
 )
 
@@ -393,7 +395,9 @@ python_unittest(
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/dialects/edge:lib",
+        "//executorch/examples/models/llama2:llama2_model",
         "//executorch/extension/pybindings:portable_lib",  # @manual
+        "//pytorch/vision:torchvision",  # @manual
     ],
 )
 
@@ -404,6 +408,7 @@ python_library(
     ],
     deps = [
         "//caffe2:torch",
+        "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
         "//executorch/exir:dim_order_utils",
         "//executorch/exir:lib",
         "//executorch/exir/capture:config",
diff --git a/exir/tests/test_memory_format_ops_pass.py b/exir/tests/test_memory_format_ops_pass.py
@@ -21,6 +21,8 @@
 )
 from executorch.exir.pass_base import ExportPass, ProxyValue
 
+from executorch.examples.models.llama2.model import Llama2Model
+
 from executorch.exir.tests.test_memory_format_ops_pass_utils import (
     MemoryFormatOpsPassTestUtils,
     MemoryFormatTestSet,
@@ -36,6 +38,8 @@
 from torch.export import export
 from torch.testing import FileCheck
 
+import torchvision
+
 
 class TestMemoryFormatOpsPass(unittest.TestCase):
     def test_op_to_copy_replacement_2d(self) -> None:
@@ -264,3 +268,67 @@ def call_operator(self, op, args, kwargs, meta):
 
         self.assertTrue(is_contiguous_dim_order(actual))
         self.assertTrue(is_contiguous_dim_order(expected))
+
+
+    def test_resnet18(self) -> None:
+        model = torchvision.models.resnet18()
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                sample_input=(torch.randn(1, 3, 224, 224),),
+                target_memory_format=torch.contiguous_format,
+                op_level_check=False,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+                atol=1e-3,
+                rtol=1e-3,
+            ),
+        )
+    
+
+    def test_resnet18_xnnpack(self) -> None:
+        model = torchvision.models.resnet18()
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                sample_input=(torch.randn(1, 3, 224, 224),),
+                target_memory_format=torch.contiguous_format,
+                op_level_check=False,
+                use_xnnpack=True,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+                atol=1e-3,
+                rtol=1e-3,
+            ),
+        )
+    
+    def test_mobilenet_v3(self) -> None:
+        model = torchvision.models.mobilenetv3.mobilenet_v3_small(pretrained=True)
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                sample_input=(torch.randn(1, 3, 224, 224),),
+                target_memory_format=torch.contiguous_format,
+                op_level_check=False,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+                atol=1e-3,
+                rtol=1e-3,
+            ),
+        )
+
+    def test_mobilenet_v3_xnnpack(self) -> None:
+        model = torchvision.models.mobilenetv3.mobilenet_v3_small(pretrained=True)
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                sample_input=(torch.randn(1, 3, 224, 224),),
+                target_memory_format=torch.contiguous_format,
+                op_level_check=False,
+                use_xnnpack=True,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+                atol=1e-3,
+                rtol=1e-3,
+            ),
+        )
diff --git a/exir/tests/test_memory_format_ops_pass_aten.py b/exir/tests/test_memory_format_ops_pass_aten.py
@@ -20,6 +20,10 @@
     _load_for_executorch_from_buffer,
 )
 
+from executorch.examples.models.llama2.model import Llama2Model
+import torchvision
+
+
 
 class TestMemoryFormatOpsPass(unittest.TestCase):
     def test_op_to_copy_replacement_2d_aten(self) -> None:
@@ -77,3 +81,34 @@ def test_op_dim_order_propagation_aten(self) -> None:
                 _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
             ),
         )
+
+    def test_resnet18(self) -> None:
+        model = torchvision.models.resnet18()
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                sample_input=(torch.randn(1, 3, 224, 224),),
+                target_memory_format=torch.contiguous_format,
+                op_level_check=False,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+                atol=1e-3,
+                rtol=1e-3,
+            ),
+        )
+    
+    
+    def test_mobilenet_v3(self) -> None:
+        model = torchvision.models.mobilenetv3.mobilenet_v3_small(pretrained=True)
+        MemoryFormatOpsPassTestUtils.memory_format_test_runner(
+            self,
+            MemoryFormatTestSet(
+                module=model.eval(),
+                sample_input=(torch.randn(1, 3, 224, 224),),
+                target_memory_format=torch.contiguous_format,
+                op_level_check=False,
+                _load_for_executorch_from_buffer=_load_for_executorch_from_buffer,
+                atol=1e-3,
+                rtol=1e-3,
+            ),
+        )
diff --git a/exir/tests/test_memory_format_ops_pass_utils.py b/exir/tests/test_memory_format_ops_pass_utils.py
@@ -11,14 +11,16 @@
 from typing import Any, Tuple
 
 import torch
-from executorch.exir import to_edge
+from executorch.exir import to_edge, to_edge_transform_and_lower
 from executorch.exir.capture._config import EdgeCompileConfig
 
 from executorch.exir.dim_order_utils import (
     is_channel_last_dim_order,
     is_contiguous_dim_order,
 )
 
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+
 from torch.export import export
 from torch.testing import FileCheck
 from torch.utils._pytree import tree_flatten
@@ -30,6 +32,10 @@ class MemoryFormatTestSet:
     sample_input: Tuple[Any, ...]
     target_memory_format: torch.memory_format
     _load_for_executorch_from_buffer: Any
+    op_level_check: bool = True
+    use_xnnpack: bool = False
+    rtol: float = 1e-05
+    atol: float = 1e-08
 
 
 class SimpleToCopyContiguousModule(torch.nn.Module):
@@ -63,27 +69,34 @@ class MemoryFormatOpsPassTestUtils:
     def memory_format_test_runner(
         test_class: unittest.TestCase, test_set: MemoryFormatTestSet
     ):
-        aten_op_str = "torch.ops.aten._to_copy.default"
-        edge_op_str = "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default"
-
         before = export(test_set.module, test_set.sample_input)
 
-        # check op strings before
-        FileCheck().check_count(aten_op_str, 1, exactly=True).check_not(
-            edge_op_str
-        ).run(before.graph_module.code)
+        if test_set.use_xnnpack:
+            epm = to_edge_transform_and_lower(
+                before, compile_config=EdgeCompileConfig(_skip_dim_order=False, _check_ir_validity=False), partitioner=[XnnpackPartitioner()])
+        else:
+            epm = to_edge(before, compile_config=EdgeCompileConfig(_skip_dim_order=False))
+
+        # check memory format ops, if needed
+        if test_set.op_level_check:
+            aten_op_str = "torch.ops.aten._to_copy.default"
+            edge_op_str = "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default"
+
+            # check op strings before
+            FileCheck().check_count(aten_op_str, 1, exactly=True).check_not(
+                edge_op_str
+            ).run(before.graph_module.code)
 
-        epm = to_edge(before, compile_config=EdgeCompileConfig(_skip_dim_order=False))
+            # check op strings
+            FileCheck().check_not(aten_op_str).check_count(
+                edge_op_str, 1, exactly=True
+            ).run(epm.exported_program().graph_module.code)
 
-        # check op strings
-        FileCheck().check_not(aten_op_str).check_count(
-            edge_op_str, 1, exactly=True
-        ).run(epm.exported_program().graph_module.code)
 
         # check EdgeOp and the new BackendOp should behave the same
         expected = before.module()(*test_set.sample_input)
         actual = epm.exported_program().module()(*test_set.sample_input)
-        test_class.assertTrue(torch.allclose(actual, expected))
+        test_class.assertTrue(torch.allclose(actual, expected, atol=test_set.atol, rtol=test_set.rtol))
         test_class.assertEqual(
             is_channel_last_dim_order(actual),
             is_channel_last_dim_order(expected),
@@ -105,7 +118,7 @@ def memory_format_test_runner(
         runtime_output = executorch_module.run_method(
             "forward", tuple(inputs_flattened)
         )[0]
-        test_class.assertTrue(torch.allclose(runtime_output, expected))
+        test_class.assertTrue(torch.allclose(runtime_output, expected,atol=test_set.atol, rtol=test_set.rtol))
         test_class.assertEqual(
             is_channel_last_dim_order(runtime_output),
             is_channel_last_dim_order(expected),
