[fx2trt] dispatch tracer improvement (#64)

Wei Wei · Wei Wei · commit 86e6e475e3cd · 2022-06-03T17:54:13.000-07:00
Summary: Pull Request resolved: https://github.com/pytorch/fx2trt/pull/64 as titled 1. able to handle in-place operation 2. handle corner case where module=ReLU and only contains aten relu with in-place=True 3. introduce functionalize to deinplace those `inplace` and `view` operations 4. introduce two methods of setting leaf nodes. leaf(op registeration) and leaf(override call_module). Only former one works with functionalize 5. More notes are in `test_dispatch_tracer.py` Reviewed By: yinghai Differential Revision: D35684102 fbshipit-source-id: ff2439587d563c4eee3a07b1ae3c1082af2679c8
diff --git a/test/tracer/test_dispatch_tracer.py b/test/tracer/test_dispatch_tracer.py
@@ -1,31 +1,140 @@
 import unittest
 
 import torch
+import torchdynamo
+import torchvision
+
+from functorch import make_fx as make_fx_pk
+from functorch.experimental import functionalize
 from fx2trt_oss.tracer.dispatch_tracer.tracer import make_fx
+from torch.library import Library
+from torchdynamo.optimizations.normalize import normalize_ir
+from torchdynamo.optimizations.python_key import fake_signature
 
 torch.manual_seed(0)
 
+wrap_lib = Library("wrap", "DEF")
+"""
+There are two methods for setting leaf_module. leaf(op registeration) and leaf(override call_module)
+Only leaf(op registeration) can work together with functionalize.
+If you do not need funcitonalize, you can choose any of the leaf module methods.
+
+Test coverage:
+PythonkeyTracerTest.test_leaf_operator_reg: python_key tracer + functionalize + leaf(op registeration)
+
+DispatchTracerTest.test_leaf_operator_reg: dispatch tracer + functionalize + leaf(op registeration)
+DispatchTracerTest.test_leaf: dispatch tracer + leaf(override call_module)
+DispatchTracerTest.test_non_tensor_input: dispatch tracer
+DispatchTracerTest.test_resnet18: dispatch tracer
+DispatchTracerTest.test_reference_copy: dispatch tracer + functionalize
+DispatchTracerTest.test_reference_copy_torchdynamo: dispatcher tracer + torchdynamo + functionalize
+"""
+
+
+class PythonkeyTracerTest(unittest.TestCase):
+    def test_leaf_operator_reg(self):
+        class Leaf(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y + torch.nn.Parameter(torch.ones(5))
+
+        leaf = Leaf()
+        wrap_lib.define("wrapped_foo(Tensor x, Tensor y) -> Tensor")
+        wrap_lib.impl("wrapped_foo", leaf, "CPU")
+
+        class Bar(torch.nn.Module):
+            def __init__(self):
+                super(Bar, self).__init__()
+                self.foo = torch.ops.wrap.wrapped_foo
+                self.other = torch.nn.Parameter(torch.ones(5))
+
+            def forward(self, x, y):
+                x = self.foo(x, y)
+                x = x + self.other
+                return x
+
+        mod = Bar()
+
+        def f(x, y):
+            return mod(x, y)
+
+        gm = make_fx_pk(functionalize(f))(torch.ones(5), torch.ones(5))
+        inputs = [torch.ones(5) + 5, torch.ones(5) + 8]
+        output = gm(*inputs)
+        ref_output = f(*inputs)
+        torch.testing.assert_close(output, ref_output)
+
 
 class DispatchTracerTest(unittest.TestCase):
-    def test_leaf_module_list(self):
-        class TestModule(torch.nn.Module):
+    def test_leaf_operator_reg(self):
+        class Leaf(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y + torch.nn.Parameter(torch.ones(5))
+
+        leaf = Leaf()
+        wrap_lib.define("wrapped_leaf(Tensor x, Tensor y) -> Tensor")
+        wrap_lib.impl("wrapped_leaf", leaf, "CPU")
+
+        class Bar(torch.nn.Module):
+            def __init__(self):
+                super(Bar, self).__init__()
+                self.leaf = torch.ops.wrap.wrapped_leaf
+                self.other = torch.nn.Parameter(torch.ones(5))
+
+            def forward(self, x, y):
+                x = self.leaf(x, y)
+                x = x + self.other
+                return x
+
+        mod = Bar()
+
+        def f(x, y):
+            return mod(x, y)
+
+        gm = make_fx(functionalize(f))(torch.ones(5), torch.ones(5))
+        inputs = [torch.ones(5) + 5, torch.ones(5) + 8]
+        output = gm(*inputs)
+        ref_output = f(*inputs)
+        torch.testing.assert_close(output, ref_output)
+        # through the op registration method, the module is defined in a call_function
+        call_function_node = None
+        for node in gm.graph.nodes:
+            if (
+                node.op == "call_function"
+                and node.target == torch.ops.wrap.wrapped_leaf
+            ):
+                call_function_node = node
+        self.assertIsNotNone(call_function_node)
+
+    def test_leaf(self):
+        class TestModuleLeaf(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.conv = torch.nn.Conv2d(3, 10, 1)
-                self.relu = torch.nn.ReLU()
+                self.relu = torch.nn.ReLU(inplace=True)
 
             def forward(self, x):
                 x = self.conv(x)
                 return self.relu(x)
 
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+                self.relu = torch.nn.ReLU(inplace=True)
+                self.leaf = TestModuleLeaf()
+
+            def forward(self, x):
+                x = self.leaf(x)
+                return self.relu(x)
+
         mod = TestModule()
 
         def f(x):
             return mod(x)
 
         a = torch.randn(1, 3, 1, 1)
         ref_output = f(a)
-        func = make_fx(f, leaf_module_list={"torch.nn.modules.activation.ReLU"})
+        func = make_fx(f, leaf_module_list={"test_dispatch_tracer.TestModuleLeaf"})
         gm = func(a)
         output = gm(a)
         torch.testing.assert_close(output, ref_output)
@@ -36,17 +145,90 @@ def f(x):
             if node.op == "call_module":
                 call_module_node = node
         self.assertIsNotNone(call_module_node)
-        self.assertEqual(call_module_node.target, "ReLU_0")
+        self.assertEqual(call_module_node.target, "TestModuleLeaf_0")
 
     def test_non_tensor_input(self):
         def foo(x):
             a = x["a"]
             b = x["b"]
             return a + b
 
-        x = {"a": torch.randn(1), "b": torch.randn(1)}
+        x = {"a": torch.randn(2, 2), "b": torch.randn(2, 2)}
         ref_output = foo(x)
         func = make_fx(foo)
         gm = func(x)
         output = gm(x)
         torch.testing.assert_close(output, ref_output)
+
+    def test_resnet18(self):
+        mod = torchvision.models.resnet18(pretrained=False)
+
+        def f(x):
+            return mod(x)
+
+        a = torch.randn(1, 3, 224, 224)
+        ref_output = f(a)
+        gm = make_fx(f)(a)
+        output = gm(a)
+        torch.testing.assert_close(output, ref_output)
+
+    def test_reference_copy(self):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x, y):
+                y[:, 0] = x[:, 0]
+                return y
+
+        mod = TestModule()
+
+        def f(x, y):
+            return mod(x, y)
+
+        a = torch.ones(2, 2) + 2
+        b = torch.ones(2, 2)
+        b_copy = torch.ones(2, 2)
+        ref_output = f(a, b)
+        gm = make_fx(functionalize(f))(a, b)
+        output = gm(a, b_copy)
+        torch.testing.assert_close(output, ref_output)
+
+    def test_reference_copy_torchdynamo(self):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.relu = torch.nn.ReLU(inplace=True)
+
+            def forward(self, x, y):
+                y = y + 3
+                y = self.relu(y)
+                y[:, 0] = x[:, 0]
+                return y
+
+        mod = TestModule()
+
+        def f(x, y):
+            return mod(x, y)
+
+        a = torch.ones(2, 2) + 2
+        b = torch.ones(2, 2)
+        inputs = [a, b]
+        ref_output = f(*inputs)
+
+        def compile_dispatch(gm, example_inputs):
+            # after normalization, relu in-place is removed
+            gm = normalize_ir(gm, example_inputs)
+            # dispatch tracer
+            nargs = len(example_inputs)
+            gm = make_fx(functionalize(fake_signature(gm, nargs)))(*example_inputs)
+            return gm
+
+        optimize_ctx = torchdynamo.optimize(
+            compile_dispatch,
+            nopython=True,
+        )
+
+        with optimize_ctx:
+            output = mod(*inputs)
+        torch.testing.assert_close(output, ref_output)
diff --git a/tracer/dispatch_tracer/tracer.py b/tracer/dispatch_tracer/tracer.py
@@ -26,17 +26,22 @@ def unwrap_proxy(e):
     return e.proxy if isinstance(e, DispatchTensor) else e
 
 
-def build_outputs(func, args, kwargs, proxy_out):
+def build_outputs(func, func_overload, args, kwargs, proxy_out, call_module=False):
     # Kind of a hacky way to test if an op is in-place or not
     if func.__name__[-1] == "_" and func.__name__[0] != "_":
         args[0].proxy = proxy_out
 
     with no_dispatch():
-        real_out = func(*args, **kwargs)
+        real_out = func_overload(*args, **kwargs)
 
     def wrap_with_proxy(e, proxy):
-        if isinstance(e, torch.Tensor):
+        if e is None:
+            e = torch.empty(())
+        if type(e) == torch.Tensor:
             return DispatchTensor(e, proxy)
+        # if module output is dispatchTensor, then all op inside it are in-place
+        elif type(e) == DispatchTensor and call_module:
+            e.proxy = proxy_out
         else:
             return e
 
@@ -46,7 +51,7 @@ def wrap_with_proxy(e, proxy):
         )
     elif isinstance(real_out, list):
         return [wrap_with_proxy(e, proxy_out[idx]) for idx, e in enumerate(real_out)]
-    elif isinstance(real_out, torch.Tensor):
+    elif type(real_out) == torch.Tensor:
         return wrap_with_proxy(real_out, proxy_out)
     else:
         return real_out
@@ -78,11 +83,12 @@ def __repr__(self):
     __torch_function__ = _disabled_torch_function_impl
 
     @classmethod
-    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
+    def __torch_dispatch__(cls, func_overload, types, args=(), kwargs=None):
+        func = func_overload.overloadpacket
         proxy_args = pytree.tree_map(unwrap_proxy, args)
         proxy_kwargs = pytree.tree_map(unwrap_proxy, kwargs)
         proxy_out = func(*proxy_args, **proxy_kwargs)
-        return build_outputs(func, args, kwargs, proxy_out)
+        return build_outputs(func, func_overload, args, kwargs, proxy_out)
 
 
 class DispatchTracer(Tracer):
@@ -101,6 +107,7 @@ def __init__(self, leaf_module_list: Optional[Set[str]] = None):
             DEFAULT_LEAF_MODULE_LIST
         )
 
+    # User can use leaf_module_list but it won't work combine with functionalize
     def call_module(
         self,
         m: torch.nn.Module,
@@ -121,7 +128,10 @@ def call_module(
             proxy_out = self.create_proxy(
                 "call_module", qualname, proxy_args, proxy_kwargs
             )
-            return build_outputs(forward, args, kwargs, proxy_out)
+
+            return build_outputs(
+                forward, forward, args, kwargs, proxy_out, call_module=True
+            )
         return forward(*args, **kwargs)
 
     def is_leaf_module(self, m) -> bool:
@@ -170,7 +180,7 @@ def dispatch_trace(
     gm = GraphModule(tracer.root, graph, name)
     gm.graph.eliminate_dead_code()
     gm.recompile()
-    return NormalizeArgs(gm).transform()
+    return gm
 
 
 def wrap_key(f, inps):
