made partitioner strip overloads

Author: Chillee

benchmarks/cse.py

@@ -3,7 +3,7 @@
 from functorch import make_fx
 from torch.profiler import profile, ProfilerActivity
 
-from functorch._src.cse import fx_graph_cse
+from functorch._src.compile_utils import fx_graph_cse
 
 def profile_it(f, inp):
     for _ in range(5):

functorch/_src/cse.py renamed to functorch/_src/compile_utils.py

@@ -1,3 +1,4 @@
+
 import torch
 import torch.fx as fx
 from torch.utils._pytree import tree_flatten
@@ -56,3 +57,16 @@ def substitute(arg_list):
                 token_map[hash_val] = token
 
     return new_graph
+
+
+def strip_overloads(gm):
+    """
+    Modifies the target of graph nodes in :attr:`gm` to strip overloads.
+
+    Args:
+        gm(fx.GraphModule): The input Fx graph module to be modified
+    """
+    for node in gm.graph.nodes:
+        if isinstance(node.target, torch._ops.OpOverload):
+            node.target = node.target.overloadpacket
+    gm.recompile()

functorch/_src/compilers.py

@@ -7,6 +7,7 @@
 from .aot_autograd import aot_function, aot_module
 from .decompositions import get_decompositions
 from .partitioners import draw_graph, min_cut_rematerialization_partition
+from .compile_utils import strip_overloads
 import time
 
 
@@ -20,19 +21,6 @@ def _canonicalize(fx_g):
     return fx_g
 
 
-def strip_overloads(gm):
-    """
-    Modifies the target of graph nodes in :attr:`gm` to strip overloads.
-
-    Args:
-        gm(fx.GraphModule): The input Fx graph module to be modified
-    """
-    for node in gm.graph.nodes:
-        if isinstance(node.target, torch._ops.OpOverload):
-            node.target = node.target.overloadpacket
-    gm.recompile()
-
-
 def ts_compile(fx_g: fx.GraphModule, _) -> Callable:
     """
     Compiles the :attr:`fx_g` with Torchscript compiler.
@@ -245,6 +233,7 @@ def nop(fx_g: fx.GraphModule, _) -> Callable:
 
 
 def simple_ts_compile(fx_g, _):
+    strip_overloads(fx_g)
     f = torch.jit.script(fx_g)
     f = torch.jit.freeze(f.eval())
     return f

functorch/_src/partitioners.py

@@ -7,7 +7,7 @@
 import os
 from torch.fx.passes import graph_drawer
 from typing import Tuple
-from .cse import fx_graph_cse
+from .compile_utils import fx_graph_cse, strip_overloads
 
 
 class InvalidNodeBase(object):
@@ -228,6 +228,7 @@ def min_cut_rematerialization_partition(
         raise RuntimeError("Need networkx installed to perform smart recomputation heuristics")
 
     #  add the CSE pass
+    strip_overloads(joint_module)
     fx_g = joint_module.graph
     cse_graph = fx_graph_cse(fx_g)
     joint_module.graph = cse_graph

test/test_memory_efficient_fusion.py

@@ -4,7 +4,7 @@
 from functorch import make_fx
 from torch.nn import functional as F
 from functorch.compile import memory_efficient_fusion
-from functorch._src.cse import fx_graph_cse
+from functorch._src.compile_utils import fx_graph_cse
 from torch.testing._internal.common_utils import TestCase, run_tests
 import inspect
 import random

test/test_pythonkey.py

@@ -496,7 +496,7 @@ def f(a, b, c, d):
         def f(x):
             return torch.mm(x, torch.ones(x.shape)).tanh().tanh()
         fw_graph, bw_graph = get_fw_bw_graph(f, [torch.randn(5, 5, requires_grad=True)])
-        self.assertEqual(get_num_ins_outs(fw_graph), (1, 2))
+        self.assertEqual(get_num_ins_outs(fw_graph), (1, 3))
 
         ins, outs = get_ins_outs(fw_graph)
         self.assertEqual(outs[1].target, torch.ops.aten.mm)