Add verifier for EXIR Aten dialect.

Summary:

Test Plan:
Unittest

Reviewers:

Subscribers:

Tasks:

Tags:

test/fx/test_verifier.py

@@ -0,0 +1,187 @@
+import unittest
+from torch.fx.verifier import (
+    ExirSpecViolationError,
+    check_valid_aten_dialect,
+    check_valid_edge_dialect,
+    check_valid_exir,
+    is_valid_aten_dialect,
+    is_valid_edge_dialect,
+    is_valid_exir,
+)
+
+
+from typing import List, Optional, Tuple, Union
+
+
+from torch.testing._internal.common_utils import TestCase
+import torch  # noqa: F401
+import torch.nn as nn
+from torch import Tensor
+import torch._dynamo as torchdynamo
+import copy
+from functorch import make_fx
+from functorch.experimental import functionalize
+
+
+@torch.no_grad()
+def capture(f, args):
+    torchdynamo.config.capture_scalar_outputs = True
+    torchdynamo.config.guard_nn_modules = True
+    torchdynamo.config.dynamic_shapes = True
+    torchdynamo.config.specialize_int_float = True
+    torchdynamo.config.allow_rnn = True
+    torchdynamo.config.verbose = True
+    torchdynamo.reset()
+    f, _ = torchdynamo.export(
+                f,
+                *copy.deepcopy(args),
+                aten_graph=True,
+                tracing_mode='fake',
+            )
+    f = functionalize(f, remove='mutations')
+    gm = make_fx(f, 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)
+
+
+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
+
+    def get_random_inputs(self) -> Tuple[Tensor, Tensor]:
+        return (torch.randn(1, 3), torch.randn(1, 3))
+
+
+class Cat(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    # def forward(self, tensors, dim=0):
+    def forward(self, *args: Tensor, dim: int) -> Tensor:
+        tensors = args[:-1]
+        return torch.cat(tensors, dim)
+
+
+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
+
+        self.layer_norm = nn.LayerNorm(input_dim)
+
+        self.relu = nn.ReLU()
+
+        self.linear1 = nn.Linear(input_dim, hidden_dim)
+        self.dropout1 = nn.Dropout()
+
+        self.linear2 = nn.Linear(hidden_dim, input_dim)
+        self.dropout2 = nn.Dropout()
+
+    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
+
+
+
+class VerifierTest(TestCase):
+
+    def test_exir_verifier(self) -> None:
+        m = ElementwiseAdd()
+        egm = capture(m, (torch.randn(100), torch.randn(100)))
+        # assert not throw
+        check_valid_exir(egm)
+        self.assertTrue(is_valid_exir(egm))
+
+    def test_exir_verifier_call_module(self) -> None:
+        m = FeedForwardBlock(10, 10)
+        gm = torch.fx.symbolic_trace(m)
+        # this would have modules that are not delegates
+        with self.assertRaises(ExirSpecViolationError):
+            check_valid_exir(gm)
+        self.assertFalse(is_valid_exir(gm))
+
+    def test_exir_verifier_no_functional(self) -> None:
+        m = ElementwiseAdd()
+        egm = capture(m, (torch.randn(100), torch.randn(100)))
+        for node in egm.graph.nodes:
+            if node.target == torch.ops.aten.add.Tensor:
+                node.target = torch.ops.aten.add.out
+        with self.assertRaises(ExirSpecViolationError):
+            check_valid_exir(egm)
+        self.assertFalse(is_valid_exir(egm))
+
+    def test_aten_dialect(self) -> None:
+        m = ElementwiseAdd()
+        egm = capture(m, (torch.randn(100), torch.randn(100)))
+        check_valid_aten_dialect(egm)
+        self.assertTrue(is_valid_aten_dialect(egm))
+
+    def test_aten_wrong_mem_format(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
+
+        m = TestModel()
+        egm = capture(m, (torch.randn(1, 3, 100, 100),))
+        egm._apply(lambda t: t.to(memory_format=torch.channels_last))
+        with self.assertRaises(ExirSpecViolationError):
+            check_valid_aten_dialect(egm)
+        self.assertFalse(is_valid_aten_dialect(egm))
+
+    def test_aten_wrong_mem_format_buffer(self) -> None:
+        class TestModel(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
+
+        m = TestModel()
+        egm = capture(m, (torch.randn(1, 3, 100, 100),))
+        egm._apply(lambda t: t.to(memory_format=torch.channels_last))
+        with self.assertRaises(ExirSpecViolationError):
+            check_valid_aten_dialect(egm)
+        self.assertFalse(is_valid_aten_dialect(egm))
+
+
+if __name__ == '__main__':
+    unittest.main()

torch/fx/verifier.py

@@ -60,7 +60,7 @@ def check_valid_exir(gm: GraphModule) -> None:  # noqa: C901
             )
             is_builtin_func = (node.target == operator.getitem or node.target.__name__ in [
                 'while_loop',
-                'cond', 
+                'cond',
             ])
             if not isinstance(node.target, OpOverload) and not is_builtin_func:
                 raise ExirSpecViolationError(
@@ -76,12 +76,13 @@ def check_valid_exir(gm: GraphModule) -> None:  # noqa: C901
             if isinstance(node.target, OpOverload):
                 stacktrace = node.meta.get("stack_trace")
 
-                if stacktrace is None:
-                    raise ExirSpecViolationError(
-                        "node of name '{}' for operator '{}' is missing stackstrace".format(
-                            node.name, op_name
-                        ),
-                    )
+                # TODO(qihqi) enable
+                # if stacktrace is None:
+                #     raise ExirSpecViolationError(
+                #         "node of name '{}' for operator '{}' is missing stackstrace".format(
+                #             node.name, op_name
+                #         ),
+                #     )
 
 
 def is_valid_exir(gm: GraphModule) -> bool:
@@ -157,4 +158,3 @@ def is_valid_edge_dialect(gm: GraphModule) -> bool:
         return True
     except ExirSpecViolationError:
         return False
-