[export] Verifier for exported program (#109519)

Summary:
X-link: pytorch/executorch#292

Added a verifier for the graph signature in a exported program

Test Plan: CI

Differential Revision: D48926643

Pull Request resolved: #109519
Approved by: https://github.com/zhxchen17

Author: angelayi

test/export/test_verifier.py

@@ -1,13 +1,8 @@
 # Owner(s): ["module: dynamo"]
-import copy
-from typing import Tuple
 import unittest
 
-import torch  # noqa: F401
-import torch.nn as nn
-import torch._dynamo as torchdynamo
-from functorch import make_fx
-from functorch.experimental import functionalize, control_flow
+import torch
+from functorch.experimental import control_flow
 from torch import Tensor
 from torch.testing._internal.common_utils import run_tests, TestCase
 from torch._dynamo.eval_frame import is_dynamo_supported
@@ -16,208 +11,218 @@
     SpecViolationError,
     Verifier,
     ATenDialectVerifier,
-)
-
-
-@torch.no_grad()
-def capture(f, args):
-    torchdynamo.config.allow_rnn = True
-    torchdynamo.reset()
-    graphmodule, _ = torchdynamo.export(
-        f,
-        *copy.deepcopy(args),
-        aten_graph=True,
-    )
-
-    def graph_with_interpreter(*args):
-        with torch.fx.traceback.preserve_node_meta():
-            return torch.fx.Interpreter(graphmodule).run(*args)
-
-    functionalized_callable = functionalize(
-        graph_with_interpreter,
-        remove='mutations_and_views',
-    )
-    gm = make_fx(functionalized_callable, tracing_mode='fake', _allow_non_fake_inputs=True)(*args)
-    return gm
-
-
-class Transpose(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def forward(self, x: Tensor, dim0: int, dim1: int) -> Tensor:
-        return x.transpose(dim0, dim1)
 
+)
+from torch._export import export
 
-class Mul(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def forward(self, input: Tensor, other: Tensor) -> Tensor:
-        # or return torch.mul(input, other)
-        return input * other
-
-    def get_random_inputs(self) -> Tuple[Tensor, Tensor]:
-        return (torch.randn(3, 2), torch.randn(3, 2))
-
-
-class ElementwiseAdd(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def forward(self, x: Tensor, y: Tensor) -> Tensor:
-        return x + y
+@unittest.skipIf(not is_dynamo_supported(), "dynamo isn't supported")
+class TestVerifier(TestCase):
+    def test_verifier_basic(self) -> None:
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            return x + y
 
-    def get_random_inputs(self) -> Tuple[Tensor, Tensor]:
-        return (torch.randn(1, 3), torch.randn(1, 3))
+        ep = export(f, (torch.randn(100), torch.randn(100)))
 
+        verifier = Verifier()
+        verifier(ep.graph_module)
 
-class Cat(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
+    def test_verifier_call_module(self) -> None:
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(10, 10)
 
-    # def forward(self, tensors, dim=0):
-    def forward(self, *args: Tensor, dim: int) -> Tensor:
-        tensors = args[:-1]
-        return torch.cat(tensors, dim)
+            def forward(self, x: Tensor) -> Tensor:
+                return self.linear(x)
 
+        gm = torch.fx.symbolic_trace(M())
 
-class FeedForwardBlock(nn.Module):
-    def __init__(self, input_dim: int, hidden_dim: int) -> None:
-        super().__init__()
-        self.input_dim = input_dim
-        self.hidden_dim = hidden_dim
+        verifier = Verifier()
+        with self.assertRaises(SpecViolationError):
+            verifier(gm)
 
-        self.layer_norm = nn.LayerNorm(input_dim)
+    def test_verifier_no_functional(self) -> None:
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            return x + y
 
-        self.relu = nn.ReLU()
+        ep = export(f, (torch.randn(100), torch.randn(100)))
+        for node in ep.graph.nodes:
+            if node.target == torch.ops.aten.add.Tensor:
+                node.target = torch.ops.aten.add_.Tensor
 
-        self.linear1 = nn.Linear(input_dim, hidden_dim)
-        self.dropout1 = nn.Dropout()
+        verifier = Verifier()
+        with self.assertRaises(SpecViolationError):
+            verifier(ep.graph_module)
 
-        self.linear2 = nn.Linear(hidden_dim, input_dim)
-        self.dropout2 = nn.Dropout()
+    def test_verifier_higher_order(self) -> None:
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            def true_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+                return x + y
 
-    def forward(self, x: Tensor) -> Tensor:
-        # LayerNorm -> Linear -> Dropout -> ReLU -> Linear -> Dropout
-        y = self.layer_norm(x)
-        y = self.linear1(y)
-        y = self.dropout1(y)
-        y = self.relu(y)
-        y = self.linear2(y)
-        y = self.dropout2(y)
-        return y
+            def false_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+                return x - y
 
-class ControlFlow(nn.Module):
+            return control_flow.cond(
+                x.shape[0] > 2, true_fn, false_fn, [x, y]
+            )
 
-    def forward(self, pred: Tensor, x: Tensor) -> Tensor:
-        return control_flow.cond(pred, lambda x: x.sin(), lambda x: x.cos(), (x,))
+        ep = export(f, (torch.randn(3, 3), torch.randn(3, 3)))
 
+        verifier = Verifier()
+        verifier(ep.graph_module)
 
-class VerifierTest(TestCase):
+    def test_verifier_nested_invalid_module(self) -> None:
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            def true_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+                return x + y
 
-    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
-    def test_verifier(self) -> None:
-        m = ElementwiseAdd()
-        egm = capture(m, (torch.randn(100), torch.randn(100)))
-        # assert not throw
-        verifier = Verifier()
-        verifier(egm)
-        self.assertTrue(verifier.is_valid(egm))
+            def false_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+                return x - y
 
-    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
-    def test_verifier_call_module(self) -> None:
-        m = FeedForwardBlock(10, 10)
-        gm = torch.fx.symbolic_trace(m)
-        # this would have modules that are not delegates
-        verifier = Verifier()
-        with self.assertRaises(SpecViolationError):
-            verifier(gm)
-        self.assertFalse(verifier.is_valid(gm))
+            return control_flow.cond(
+                x.shape[0] > 2, true_fn, false_fn, [x, y]
+            )
 
-    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
-    def test_verifier_no_functional(self) -> None:
-        m = ElementwiseAdd()
-        egm = capture(m, (torch.randn(100), torch.randn(100)))
-        for node in egm.graph.nodes:
+        ep = export(f, (torch.randn(3, 3), torch.randn(3, 3)))
+        for node in ep.graph_module.true_graph_0.graph.nodes:
             if node.target == torch.ops.aten.add.Tensor:
-                node.target = torch.ops.aten.add.out
+                node.target = torch.ops.aten.add_.Tensor
+
         verifier = Verifier()
         with self.assertRaises(SpecViolationError):
-            verifier(egm)
-        self.assertFalse(verifier.is_valid(egm))
-
-    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
-    def test_aten_dialect(self) -> None:
-        m = ElementwiseAdd()
-        egm = capture(m, (torch.randn(100), torch.randn(100)))
-        verifier = ATenDialectVerifier()
-        verifier(egm)
-        self.assertTrue(verifier.is_valid(egm))
+            verifier(ep.graph_module)
 
-    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
-    def test_aten_wrong_mem_format(self) -> None:
+    def test_aten_verifier_wrong_op(self) -> None:
         class TestModel(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.a = torch.nn.parameter.Parameter(
-                    torch.randn(1, 3, 100, 100).to(memory_format=torch.channels_last)
-                )
 
             def forward(self, x):
-                return self.a + x
+                return torch.ops.aten._add_relu(x, x)
 
         m = TestModel()
-        egm = capture(m, (torch.randn(1, 3, 100, 100),))
-        egm._apply(lambda t: t.to(memory_format=torch.channels_last))
+        egm = torch.fx.symbolic_trace(m)
         verifier = ATenDialectVerifier()
         with self.assertRaises(SpecViolationError):
             verifier(egm)
         self.assertFalse(verifier.is_valid(egm))
 
-    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
-    def test_aten_wrong_mem_format_buffer(self) -> None:
-        class TestModel(torch.nn.Module):
+    def test_ep_verifier_basic(self) -> None:
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(10, 10)
+
+            def forward(self, x: Tensor) -> Tensor:
+                return self.linear(x)
+
+        ep = export(M(), (torch.randn(10, 10),))
+        ep._validate()
+
+    def test_ep_verifier_invalid_param(self) -> None:
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            return x + y
+
+        ep = export(f, (torch.randn(100), torch.randn(100)))
+
+        # Parameter doesn't exist in the state dict
+        ep.graph_signature.parameters.append("bad_param")
+        with self.assertRaisesRegex(SpecViolationError, "not in the state dict"):
+            ep._validate()
+
+        # Add non-torch.nn.Parameter parameter to the state dict
+        ep.state_dict["bad_param"] = torch.randn(100)
+        with self.assertRaisesRegex(
+            SpecViolationError, "not an instance of torch.nn.Parameter"
+        ):
+            ep._validate()
+
+        # Add torch.nn.Parameter to state dict, but this should still error
+        # because there are an incorrect number of placeholder nodes
+        ep.state_dict["bad_param"] = torch.nn.Parameter(torch.randn(100))
+        with self.assertRaisesRegex(
+            SpecViolationError, "not found in the exported program's parameter list"
+        ):
+            ep._validate()
+
+    def test_ep_verifier_invalid_buffer(self) -> None:
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            return x + y
+
+        ep = export(f, (torch.randn(100), torch.randn(100)))
+
+        # Buffer doesn't exist in the state dict
+        ep.graph_signature.buffers.append("bad_buffer")
+        with self.assertRaisesRegex(SpecViolationError, "not in the state dict"):
+            ep._validate()
+
+        # Incorrect number of placeholder nodes
+        ep.state_dict["bad_buffer"] = torch.randn(100)
+        with self.assertRaisesRegex(
+            SpecViolationError, "not found in the exported program's buffer list"
+        ):
+            ep._validate()
+
+    def test_ep_verifier_buffer_mutate(self) -> None:
+        class M(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.register_buffer(
-                    "a",
-                    torch.randn(1, 3, 100, 100).to(memory_format=torch.channels_last),
-                )
 
-            def forward(self, x):
-                return self.a + x
+                self.my_parameter = torch.nn.Parameter(torch.tensor(2.0))
 
-        m = TestModel()
-        egm = capture(m, (torch.randn(1, 3, 100, 100),))
-        egm._apply(lambda t: t.to(memory_format=torch.channels_last))
-        verifier = ATenDialectVerifier()
-        with self.assertRaises(SpecViolationError):
-            verifier(egm)
-        self.assertFalse(verifier.is_valid(egm))
+                self.register_buffer("my_buffer1", torch.tensor(3.0))
+                self.register_buffer("my_buffer2", torch.tensor(4.0))
 
-    def test_aten_wrong_op(self) -> None:
-        class TestModel(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
+            def forward(self, x1, x2):
+                # Use the parameter, buffers, and both inputs in the forward method
+                output = (
+                    x1 + self.my_parameter
+                ) * self.my_buffer1 + x2 * self.my_buffer2
 
-            def forward(self, x):
-                return torch.ops.aten._add_relu(x, x)
+                # Mutate one of the buffers (e.g., increment it by 1)
+                self.my_buffer2.add_(1.0)
+                return output
 
-        m = TestModel()
-        egm = torch.fx.symbolic_trace(m)
-        verifier = ATenDialectVerifier()
-        with self.assertRaises(SpecViolationError):
-            verifier(egm)
-        self.assertFalse(verifier.is_valid(egm))
+        ep = export(M(), (torch.tensor(5.0), torch.tensor(6.0)))
+        ep._validate()
 
-    @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported")
-    def test_verifier_control_flow_success(self) -> None:
-        m = ControlFlow()
-        gm = torch._export.export(m, (torch.tensor(True), torch.randn(3, 4))).graph_module
-        # No error should be raised
-        verifier = ATenDialectVerifier()
-        verifier(gm)
+    def test_ep_verifier_invalid_output(self) -> None:
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+                self.my_parameter = torch.nn.Parameter(torch.tensor(2.0))
+
+                self.register_buffer("my_buffer1", torch.tensor(3.0))
+                self.register_buffer("my_buffer2", torch.tensor(4.0))
+
+            def forward(self, x1, x2):
+                # Use the parameter, buffers, and both inputs in the forward method
+                output = (
+                    x1 + self.my_parameter
+                ) * self.my_buffer1 + x2 * self.my_buffer2
+
+                # Mutate one of the buffers (e.g., increment it by 1)
+                self.my_buffer2.add_(1.0)
+                return output
+
+        ep = export(M(), (torch.tensor(5.0), torch.tensor(6.0)))
+
+        output_node = list(ep.graph.nodes)[-1]
+        with ep.graph.inserting_before(output_node):
+            additional_output_node = ep.graph.call_function(
+                torch.add, args=(output_node.args[0][0], output_node.args[0][0])
+            )
+            output_node.args = (
+                (
+                    output_node.args[0][0],
+                    additional_output_node,
+                    output_node.args[0][1],
+                ),
+            )
+
+        with self.assertRaisesRegex(SpecViolationError, "Number of output nodes"):
+            ep._validate()
 
 
 if __name__ == '__main__':