Use Functionalization pass (#810)

Author: anijain2305

functorch/_src/aot_autograd.py

@@ -8,6 +8,8 @@
 import torch.utils.dlpack
 from torch.nn.utils import _stateless
 from functorch._C import CompileCache
+from functorch.experimental import functionalize
+from . import config
 from .decompositions import register_decomposition
 from .partitioners import default_partition
 from .named_members_polyfill import _named_parameters, _named_buffers
@@ -188,6 +190,12 @@ def forward(ctx, *flat_tensor_args):
                             *joint_inputs
                         )
 
+                        if config.use_functionalize:
+                            # Functionalize the foward backward graph. First create a
+                            # fake fn to make functionalize happy
+                            def fake_fn(primals, tangents):
+                                return fx_g(primals, tangents)
+                            fx_g = make_fx(functionalize(fake_fn))(*joint_inputs)
                 fw_module, bw_module = partition_fn(fx_g, joint_inputs)
                 # print(fw_module.code, bw_module.code)
 

functorch/_src/config.py

@@ -0,0 +1,11 @@
+# Copyright (c) Facebook, Inc. and its 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.
+
+"""
+Global flags for aot autograd
+"""
+
+use_functionalize = False

functorch/compile/__init__.py

@@ -27,3 +27,4 @@
     draw_graph,
     draw_joint_graph,
 )
+from .._src import config

test/test_functionalize.py

@@ -0,0 +1,52 @@
+import torch
+
+import functorch
+from torch.testing._internal.common_utils import run_tests, TestCase, IS_WINDOWS
+import unittest
+from unittest.mock import patch
+import functools
+
+from functorch.compile import aot_function, nop
+import test_compile_cache
+import test_pythonkey
+
+
+def make_functionalize_fn(fn):
+    @functools.wraps(fn)
+    def _fn(*args, **kwargs):
+        with patch.object(functorch.compile.config, "use_functionalize", True):
+            return fn(*args, **kwargs)
+
+    return _fn
+
+
+def make_functionalize_test(cls):
+    class FunctionalizeTest(cls):
+        pass
+
+    FunctionalizeTest.__name__ = f"Functionalize{cls.__name__}"
+
+    for name in dir(cls):
+        if name.startswith("test_"):
+            fn = getattr(cls, name)
+            if not callable(fn):
+                continue
+
+            new_name = f"{name}_functionalize"
+            fn = make_functionalize_fn(fn)
+            fn.__name__ = new_name
+            setattr(FunctionalizeTest, name, None)
+            setattr(FunctionalizeTest, new_name, fn)
+
+    return FunctionalizeTest
+
+
+FunctionalizeTestCompileCache = make_functionalize_test(test_compile_cache.TestCompileCache) 
+FunctionalizeTestCompileCacheStaticArgs = make_functionalize_test(test_compile_cache.TestCompileCacheStaticArgs)
+FunctionalizeTestPythonKeyAOT = make_functionalize_test(test_pythonkey.TestAOTAutograd)
+FunctionalizeTestPythonKeyContiguous = make_functionalize_test(test_pythonkey.TestContiguous)
+FunctionalizeTestPythonKeyRandom = make_functionalize_test(test_pythonkey.TestRandom)
+FunctionalizeTestPythonKeyPartitioning = make_functionalize_test(test_pythonkey.TestPartitioning)
+
+if __name__ == "__main__":
+    run_tests()