Did some dead code elimination before partitioning

Author: Chillee

functorch/_src/compilers.py

@@ -266,6 +266,7 @@ def nnc_jit(f, static_argnums=None):
     aten.hardswish,
     aten.hardsigmoid,
 }
+
 default_decompositions = get_decompositions(default_decompositions)
 
 

functorch/_src/partitioners.py

@@ -227,9 +227,12 @@ def min_cut_rematerialization_partition(
     except ImportError:
         raise RuntimeError("Need networkx installed to perform smart recomputation heuristics")
 
-    #  add the CSE pass
     strip_overloads(joint_module)
+    joint_module.graph.eliminate_dead_code()
+    joint_module.recompile()
     fx_g = joint_module.graph
+
+    #  add the CSE pass
     cse_graph = fx_graph_cse(fx_g)
     joint_module.graph = cse_graph
     full_bw_graph = joint_module.graph
@@ -373,8 +376,6 @@ def get_node_weight(node):
     # To make this stuff deterministic
     node_idx = {node: idx for idx, node in enumerate(joint_module.graph.nodes)}
     saved_values = sorted((name_to_node[node] for node in cut_nodes), key=lambda x: node_idx[x])
-    saved_values = [name_to_node[node] for node in cut_nodes]
-
     return _extract_fwd_bwd_modules(joint_module, saved_values)
 
 

test/test_pythonkey.py

@@ -314,7 +314,6 @@ class TestEagerFusionOpInfo(TestCase):
         xfail('linalg.cholesky'),
         skip('msort'),
         xfail('nn.functional.dropout'),
-        xfail('polar'),
         xfail('to_sparse'),
         xfail('addcdiv'),
         xfail('cholesky'),
@@ -326,6 +325,8 @@ class TestEagerFusionOpInfo(TestCase):
         xfail('matrix_exp'),
         xfail('trapezoid'),
         xfail('trapz'),
+        xfail('corrcoef'),
+        xfail('cov'),
         skip('nn.functional.binary_cross_entropy_with_logits'),  # seems to fail sometimes?
         skip('nn.functional.margin_ranking_loss'),  # seems flaky
     })