added some tests and stuff to minimizer

Author: Chillee

functorch/_src/fx_minifier.py

@@ -42,62 +42,60 @@ def _convert_node_to_placeholder(node, inps):
         for tuple_user in list(node.users):
             _convert_node_to_placeholder(tuple_user, inps)
 
-def minimizer(fail_f: fx.GraphModule, inps, pass_checker):
+def minimizer(fail_f: fx.GraphModule, inps, module_fails):
     """
-    Minimizes a FX graph with given inputs, such that the resulting FX graph still fails the pass_checker.
+    Minimizes a FX graph with given inputs, such that the resulting FX graph still returns True for module_fails.
 
     Does 2 main strategies:
     1. Truncates suffix: Removes some suffix from the graph and sets a new output.
     2. Delta Debugging: Tries replacing half of the graph with inputs. If fails, tries replacing quarter of the graph, etc.
 
     >>> failing_function = fx.symbolic_trace(f)
     >>> minimize(failing_function, [torch.randn(5)], lambda fx_g, inps: fx_g(*inps))
+
+    note: module_fails returns True if it fails.
     """
     failing_graph = fail_f.graph
     cur_size = len(failing_graph.nodes)
 
-    def graph_passes(graph, inps):
-        graph.lint()
+    def graph_fails(graph, inps):
         mod = fx.GraphModule(fail_f, graph)
-        return pass_checker(mod, inps)
+        mod.graph.lint()
+        return module_fails(mod, inps)
 
     ConcreteProp(fail_f).propagate(*inps)
-    if graph_passes(failing_graph, inps):
+    if not graph_fails(failing_graph, inps):
         raise RuntimeError("Input graph did not fail the tester")
     print(f"Started off with {cur_size} nodes")
 
     def remove_suffix(cur_graph, cur_inps):
         print("Strategy: Remove suffix")
-        assert not graph_passes(cur_graph, cur_inps)
+        assert graph_fails(cur_graph, cur_inps)
         gap = 2**math.floor(math.log2(len(cur_graph.nodes)))
         tested = set()
         while gap >= 1:
-            print(f"search gap: {gap}: ", end='')
             new_graph = fx.Graph()
             env = {}
             for idx, node in enumerate(cur_graph.nodes):
                 new_node = new_graph.node_copy(node, lambda x: env[x])
                 if node.op not in ['placeholder', 'output']:
                     if idx % gap == 0 and idx not in tested:
-                        print(f"{idx}", end=',')
-                        output_node = new_graph.output(([new_node],))
-                        if not graph_passes(new_graph, cur_inps) and len(new_graph.nodes) < len(cur_graph.nodes):
+                        output_node = new_graph.output((new_node,))
+                        if graph_fails(new_graph, cur_inps) and len(new_graph.nodes) < len(cur_graph.nodes):
                             print()
-                            print(f"SUCCESS: Found failing case with first {idx} nodes")
+                            print(f"SUCCESS: Removed [{idx}:{len(cur_graph.nodes)})")
                             return (new_graph, cur_inps), True
                         else:
                             tested.add(idx)
                             new_graph.erase_node(output_node)
                 env[node] = new_node
             gap //= 2
-            print()
         print("FAIL: Could not remove suffix")
         return (cur_graph, cur_inps), False
 
 
     def remove_unused_inputs(cur_graph, cur_inps):
-        print("Strategy: Remove unused inputs")
-        assert not graph_passes(cur_graph, cur_inps)
+        assert graph_fails(cur_graph, cur_inps)
         ph_nodes = _get_placeholders(cur_graph)
         if len(ph_nodes) != len(cur_inps):
             print(cur_graph)
@@ -111,13 +109,22 @@ def remove_unused_inputs(cur_graph, cur_inps):
             else:
                 new_inps.append(cur_inps[idx])
 
-        if len(new_inps) < len(cur_inps):
+        if len(new_inps) < len(cur_inps) and graph_fails(cur_graph, new_inps):
+            print("Strategy: Remove unused inputs")
             print(f"SUCCESS: Went from {len(cur_inps)} inputs to {len(new_inps)} inputs")
             return (cur_graph, new_inps), True
         else:
-            print("FAIL: Could not remove inputs")
             return (cur_graph, new_inps), False
 
+    def eliminate_dead_code(cur_graph, cur_inps):
+        orig_size = len(cur_graph.nodes)
+        if cur_graph.eliminate_dead_code() and graph_fails(cur_graph, cur_inps):
+            print("Strategy: Eliminate dead code")
+            print(f"SUCCESS: Went from {orig_size} nodes to {len(cur_graph.nodes)} nodes")
+            return (cur_graph, cur_inps), True
+        else:
+            return (cur_graph, cur_inps), False
+
     def consolidate_placeholders(cur_graph):
         new_graph = fx.Graph()
         env = {}
@@ -134,47 +141,49 @@ def consolidate_placeholders(cur_graph):
 
     def delta_debugging(cur_graph: fx.Graph, cur_inps):
         print("Strategy: Delta Debugging")
-        assert not graph_passes(cur_graph, cur_inps)
+        assert graph_fails(cur_graph, cur_inps)
         starting_placeholders = len(_get_placeholders(cur_graph))
         num_nodes = len(cur_graph.nodes)
         gap = int(2**math.floor(math.log2(num_nodes)))
         while gap >= 1:
-            print(f"Searching with gap of {gap}")
             for start_range in range(0, num_nodes, gap):
                 is_removing = False
                 new_graph = copy.deepcopy(cur_graph)
                 new_inps = cur_inps[:]
-                for idx in range(start_range, min(num_nodes, start_range + gap)):
+                end_range = min(num_nodes, start_range + gap)
+                for idx in range(start_range, end_range):
                     new_node = list(new_graph.nodes)[idx]
                     if new_node.op not in ['placeholder', 'output']:
                         is_removing = True
                         _convert_node_to_placeholder(new_node, new_inps)
                 if not is_removing:
                     continue 
                 new_graph = consolidate_placeholders(new_graph)
-                if not graph_passes(new_graph, new_inps):
-                    print(f"SUCCESS: Went from {starting_placeholders} placeholders to {len(_get_placeholders(new_graph))}")
+                if graph_fails(new_graph, new_inps):
+                    print(f"SUCCESS: Removed ({start_range}:{end_range}] - Went from {starting_placeholders} placeholders to {len(_get_placeholders(new_graph))}")
                     return (new_graph, new_inps), True
             gap //= 2
 
         print("FAIL: Could not remove prefix")
         return (cur_graph, inps), False
 
 
+    print(f"###################")
+    print(f"Current size: {len(failing_graph.nodes)}")
+    print(f"###################")
     while True:
         any_succeeded = False
-        for strategy in [remove_suffix, remove_unused_inputs, delta_debugging, remove_unused_inputs]:
-            print(f"###################")
-            print(f"Current size: {len(failing_graph.nodes)}")
-            print(f"###################")
+        for strategy in [remove_suffix, eliminate_dead_code, remove_unused_inputs, delta_debugging, eliminate_dead_code, remove_unused_inputs]:
             out = strategy(copy.deepcopy(failing_graph), inps[:])
             (cur_graph, cur_inps), succeeded = out
             if succeeded:
+                print()
+                print(f"###################")
+                print(f"Current size: {len(cur_graph.nodes)}")
+                print(f"###################")
                 failing_graph = cur_graph
-                failing_graph.eliminate_dead_code()
                 inps = cur_inps
                 any_succeeded = True
-            print()
 
         if not any_succeeded:
             break

test/test_minifier.py

@@ -0,0 +1,53 @@
+import torch
+from torch.nn import functional as F
+from functorch.compile import minimizer
+from functorch import make_fx
+from torch.testing._internal.common_utils import TestCase, run_tests
+from typing import Callable
+
+
+class TestMinifier(TestCase):
+    def test_has_mul_minifier(self):
+        def failing_f(x, y):
+            y = y / 3
+            x = x + 3
+            x = x * y
+            return x + y
+        inps = [torch.randn(3), torch.randn(3)]
+        failing_f = make_fx(failing_f)(*inps)
+
+        def pass_checker(fx_g, inps):
+            return (torch.ops.aten.mul in set([i.target for i in fx_g.graph.nodes]))
+
+        min_f, inps = minimizer(failing_f, inps, pass_checker)
+        assert len(min_f.graph.nodes) == 4
+        assert len(inps) == 2 
+
+    def test_has_add_mul(self):
+        def failing_f(x):
+            x = x * 3
+            x = x + 5
+            x = x.cos()
+            zero = x - x
+            result = zero / zero
+            result = result + 3
+            return (result * 2,)
+
+        inps = [torch.randn(3)]
+        failing_f = make_fx(failing_f)(*inps)
+
+        def pass_checker(fx_g, inps):
+            # Basically, make sure none of the inputs are nans
+            for i in inps:
+                if torch.isnan(i).any():
+                    return False
+            return torch.isnan(fx_g(*inps)[0]).any()
+
+        min_f, inps = minimizer(failing_f, inps, pass_checker)
+        import pdb; pdb.set_trace()
+        assert len(min_f.graph.nodes) == 3
+        assert len(inps) == 1 
+
+
+if __name__ == "__main__":
+    run_tests()