Adds save_thunderfx_repros to save scripts for all the subgraphs and optionally save fusion region and traces (#2232)

While exploring the split reasons and observing changes in traces and fusion areas when adding missing fallback operators, I realized that having an API to save all repro scripts, with options to also save fusion graphs or traces to files, would be helpful. Before this PR, the user had to do the iteration and save manually.

Author: kiya00

thunder/dynamo/report.py

@@ -10,6 +10,7 @@
 import copy
 from itertools import chain
 from looseversion import LooseVersion
+import shutil
 
 import torch
 from thunder.core.pytree import tree_flatten
@@ -1394,3 +1395,111 @@ def save_failing_repros(
             report.write_repro(
                 repros_folder, compile_fn, extra_comment_str=comment, check_consistency=check_consistency
             )
+
+
+def create_folder(folder_path: str | PathLike, force_overwrite: bool = False):
+    folder_path = Path(folder_path)
+
+    if folder_path.exists():
+        if not folder_path.is_dir():
+            raise RuntimeError(f"{folder_path} exists and is not a directory.")
+
+        if force_overwrite:
+            shutil.rmtree(folder_path)
+        else:
+            raise RuntimeError(f"Folder {folder_path} already exists. Use force_overwrite=True to overwrite.")
+
+    folder_path.mkdir(parents=True, exist_ok=False)
+
+
+def save_thunderfx_repros(
+    fn: Callable,
+    folder_path: str | PathLike,
+    *,
+    use_benchmark: bool = False,
+    check_runnability: bool = False,
+    save_fusion: bool = False,
+    save_trace: bool = False,
+    stream: TextIO = sys.stdout,
+    force_overwrite: bool = False,
+    **compile_kwargs,
+):
+    """
+    Saves reproduction scripts for ThunderFX subgraphs.
+
+    This function:
+    1. Creates a folder structure to organize the repros
+    .
+    └── graph0
+        ├── fusion_reports
+        │   ├── graph0_thunder_0_nvFusion0_forward_repro_nvfuser.py
+        │   ├── graph0_thunder_0_nvFusion1_forward_repro_nvfuser.py
+        │   ├── graph0_thunder_0_nvFusion2_backward_repro_nvfuser.py
+        ├── graph0_thunder_0_bwd_trace.py
+        ├── graph0_thunder_0_fwd_trace.py
+        └── graph0_thunder_0.py
+
+    2. For each Thunder FX graph and its subgraphs:
+        - Checks runnability if requested
+        - Saves benchmark or repro scripts
+        - Saves trace information if requested
+        - Saves nvFusion repros if requested
+
+    Args:
+        fn: The callable to analyze
+        folder_path: Path to save repros to
+        use_benchmark: If True, saves benchmark scripts instead of repros
+        check_runnability: If True, checks if graphs can run with Thunder
+        save_fusion: If True, saves nvFusion repros
+        save_trace: If True, saves trace information
+        stream: Stream to write output log informationto
+        force_overwrite: If True, overwrites existing folder at folder_path
+        **compile_kwargs: Keyword arguments for Thunder and torch.compile
+
+    Returns:
+        A wrapped function that saves repros when called with inputs
+    """
+    from thunder.dynamo.utils import get_torch_compile_kwargs
+
+    folder_path = Path(folder_path)
+    create_folder(folder_path, force_overwrite)
+    torch_compile_kwargs = get_torch_compile_kwargs(**compile_kwargs)
+    thunder_jit_kwargs = {k: v for k, v in compile_kwargs.items() if k not in torch_compile_kwargs}
+    thunderjit = ThunderCompileSpecification(**thunder_jit_kwargs)
+
+    def inner_fn(*args, **kwargs):
+        thunder_fxgraph_reports = get_thunder_fxgraph_reports(fn, stream=stream, **compile_kwargs)(*args, **kwargs)
+        for thunder_fxgraph_report in thunder_fxgraph_reports:
+            graph_folder = folder_path / thunder_fxgraph_report.graph_name
+            graph_folder.mkdir(exist_ok=True, parents=True)
+            for split_report in thunder_fxgraph_report.subgraph_reports:
+                if check_runnability or save_trace or save_fusion:
+                    try:
+                        split_report.create_fusion_reports()
+                    except Exception as e:
+                        stream.write(f"Failed to run the {split_report.graph_name} using Thunder with exception: {e}\n")
+                        split_report.write_repro(
+                            graph_folder, thunderjit, file_name=f"failed_{split_report.graph_name}.py"
+                        )
+                        continue
+                    else:
+                        stream.write(f"Successfully ran the {split_report.graph_name} using Thunder\n")
+                if use_benchmark:
+                    split_report.write_benchmark(graph_folder, thunderjit, WallTime)
+                else:
+                    split_report.write_repro(graph_folder, thunderjit)
+                if save_trace:
+                    with open(graph_folder / f"{split_report.graph_name}_fwd_trace.py", "w") as f:
+                        f.write(str(split_report.fwd_trc))
+                    with open(graph_folder / f"{split_report.graph_name}_bwd_trace.py", "w") as f:
+                        f.write(str(split_report.bwd_trc))
+                if save_fusion:
+                    fusion_folder = graph_folder / "fusion_reports"
+                    fusion_folder.mkdir(exist_ok=True, parents=True)
+                    for fusion_report in split_report.fusion_reports:
+                        if use_benchmark:
+                            fusion_report.write_nvfuser_benchmark(fusion_folder, WallTime)
+                        else:
+                            fusion_report.write_nvfuser_repro(fusion_folder)
+
+    return inner_fn