Naming

Author: anijain2305

torchdynamo/config.py

@@ -99,18 +99,16 @@
     "torch._decomp",
 }
 
+# Debug Flag to try minifier at different stages. Possible values are {None, "aot", "dynamo"}
+# None - Minifier is switched off
+# dynamo - Runs minifier on the TorchDynamo produced graphs, if compilation fails
+# aot - Runs minifier on the Aot Autograd produced graphs, if compilation fails
+repro_after = os.environ.get("REPRO_AFTER", None)
 # Compiler compilation debug info
-# 0: Nothing printed out when compilation fails
-# 1: Dump the graph out to repro.py if compilation fails
-# 2: Dumps the graph out to minify_repro.py with a minifier if compilation fails
-# 3: Always dumps the last graph ran out to minify_repro.py, useful for segfaults/irrecoverable errors
-repro_level = int(os.environ.get("COMPILER_REPRO_LEVEL", 0))
-
-# TorchDynamo produced Fx GraphModule compilation debug info
-# 0: Nothing printed out when compilation fails
-# 1: Dump the initial graph to repro.tar.gz
-# 2/3: Minifies and Dumps the minified repro to repro.tar.gz
-backend_repro_level = int(os.environ.get("BACKEND_REPRO_LEVEL", 0))
+# 1: Dumps the original graph out to repro.py/repro.tar.gz if compilation fails
+# 2: Either 1) Dumps the original graph out to minify_repro.py with a minifier if compilation fails
+#        or 2) Minifies and dumps the minified graph to repro.py/repro.tar.gz if compilation fails
+repro_level = int(os.environ.get("REPRO_LEVEL", 2))
 
 # Not all backends support scalars. Some calls on torch.Tensor (like .item()) return a scalar type.
 # When this flag is set to False, we introduce a graph break instead of capturing.

torchdynamo/debug_utils.py

@@ -310,25 +310,23 @@ def wrap_compiler_debug(compiler, compiler_name: str):
     @functools.wraps(compiler)
     def debug_wrapper(gm, example_inputs, **kwargs):
         orig_graph = copy.deepcopy(gm.graph)
-        if config.repro_level == 3:
-            dump_to_minify(
-                fx.GraphModule(gm, orig_graph), example_inputs, compiler_name
-            )
-
-        try:
-            compiled_fn = compiler(gm, example_inputs, **kwargs)
-            if config.repro_level > 0:
+        assert config.repro_after in ("dynamo", "aot", None)
+        if config.repro_after == "aot":
+            try:
+                compiled_fn = compiler(gm, example_inputs, **kwargs)
                 compiled_fn(*example_inputs)
-        except Exception as e:
-            if config.repro_level == 1:
-                dump_compiler_graph_state(
-                    fx.GraphModule(gm, orig_graph), example_inputs, compiler_name
-                )
-            elif config.repro_level == 2:
-                dump_to_minify(
-                    fx.GraphModule(gm, orig_graph), example_inputs, compiler_name
-                )
-            raise e
+            except Exception as e:
+                if config.repro_level == 1:
+                    dump_compiler_graph_state(
+                        fx.GraphModule(gm, orig_graph), example_inputs, compiler_name
+                    )
+                elif config.repro_level == 2:
+                    dump_to_minify(
+                        fx.GraphModule(gm, orig_graph), example_inputs, compiler_name
+                    )
+                raise e
+        else:
+            compiled_fn = compiler(gm, example_inputs, **kwargs)
 
         return compiled_fn
 
@@ -517,11 +515,12 @@ def wrap_backend_debug(compiler_fn, compiler_name: str):
 
     @functools.wraps(compiler_fn)
     def debug_wrapper(gm, example_inputs, **kwargs):
-        compiled_gm = compiler_fn(gm, example_inputs, **kwargs)
-        if config.backend_repro_level > 0:
+        assert config.repro_after in ("dynamo", "aot", None)
+        if config.repro_after == "dynamo":
             # Ensure that we fail when backend fails
             config.raise_on_backend_error = True
             try:
+                compiled_gm = compiler_fn(gm, example_inputs, **kwargs)
                 run_fwd_maybe_bwd(compiled_gm, clone_inputs(example_inputs))
             except Exception as exc:
                 orig_failure = str(exc)
@@ -531,7 +530,7 @@ def debug_wrapper(gm, example_inputs, **kwargs):
                 dump_state_fn = functools.partial(
                     dump_backend_state, compiler_name=compiler_name
                 )
-                if config.backend_repro_level == 1:
+                if config.repro_level == 1:
                     dump_state_fn(
                         fx.GraphModule(gm, copy.deepcopy(gm.graph)), example_inputs
                     )
@@ -558,6 +557,8 @@ def debug_wrapper(gm, example_inputs, **kwargs):
                         dump_state=dump_state_fn,
                     )
                     raise exc
+        else:
+            compiled_gm = compiler_fn(gm, example_inputs, **kwargs)
 
         return compiled_gm