Repository: Lightning-AI/lightning-thunder
Issue category: Feature Request

Issue:
[Reporting tool] Adds the Thunder segmentation and fusion information to the FXReport
There are two immediate next options I can think of:

One option is to make thunder segmentation and fusion information programmatically available. This could be done by having a function like `analyze_with_thunder` that takes an FXGraphReport and creates a ThunderFXGraphReport, which would be a subclass of FXGraphReport. The ThunderFXGraphReport would have additional ThunderSegmentReport classes, and maybe ThunderFusionReport classes, so that the segments and fusions could be iterated over. Then the ability to write reproducers from the segments alone and from the fusions alone could be added to those classes (reusing the existing functionality as much as possible). That would allow for the creation of files that reproduce or benchmark segments and fusions. Each nvFusion could then be compared to executing with torch.compile instead, for example.

I think the second option is to focus on making the timing data programmatically available at the FX level. Basically running the benchmark and returning its data as a Python object. In the future we should let callers specify their own timers, but a first PR can just pick a timer to use. 

Whichever is done first, the second one can probably be done next. The timings don't need to be perfect at first, but it will be exciting to work on the proper timing of fx graphs, segments, and fusions in the near future. @kshitij12345 did some work attempting to compare nvfuser and torch.compile fusions, but the nvFuser team and mixology teams may have some other ideas for timing, too.

Once we have programmatically available timings for fx graphs, segments, and fusions, then I think it's time to revisit running these reports as part of the reports that Mixology produces.

How does that sound, @kiya00?

_Originally posted by @mruberry in https://github.com/Lightning-AI/lightning-thunder/pull/1707#pullrequestreview-2587842277_
            

cc @crcrpar


Hints:

Patch:
diff --git a/thunder/core/utils.py b/thunder/core/utils.py
index dc4e6f345d..cc5c5995cc 100644
--- a/thunder/core/utils.py
+++ b/thunder/core/utils.py
@@ -17,7 +17,7 @@
 from thunder.core.proxies import Proxy, NumberProxy, TensorProxy, variableify, CONSTRAINT, Variable
 from thunder.core.baseutils import *
 from thunder.core.codeutils import *
-from thunder.core.trace import TraceCtx
+from thunder.core.trace import TraceCtx, tracectx
 import thunder.core.prims as prims
 
 if TYPE_CHECKING:
@@ -1246,3 +1246,16 @@ def get_dtype_for_device_if_enabled(self, query_device_str: str):
 
         # Not found on the stack.
         return None
+
+
+def create_python_callable_from_bsym(bsym: BoundSymbolInterface) -> str:
+    trace = TraceCtx()
+    si = SigInfo(bsym.sym.name)
+    si.args = [(v.name, None) for v in bsym.flat_args]
+    trace._siginfo = si
+    trace.bound_symbols = list(bsym.subsymbols)
+
+    with tracectx(trace):
+        prims.python_return(bsym.output)
+
+    return trace.python(include_decorators=False)
diff --git a/thunder/dynamo/report.py b/thunder/dynamo/report.py
index 34d55edcea..8ebb265f2d 100644
--- a/thunder/dynamo/report.py
+++ b/thunder/dynamo/report.py
@@ -10,10 +10,22 @@
 
 import torch
 from thunder.core.pytree import tree_flatten
-from thunder.core.utils import sequencify
+from thunder.core.utils import sequencify, create_python_callable_from_bsym
 from thunder.dynamo.compiler import thunderfx
-from thunder.dynamo.utils import _get_example_inputs_from_placeholder, _readable, arg_like, get_env
-from thunder.dynamo.repro_script_template import benchmark_multi_exe_code_template, repro_code_template
+from thunder.dynamo.utils import (
+    _get_example_inputs_from_placeholder,
+    _readable,
+    arg_like,
+    get_env,
+    get_split_reasons_string,
+    CompilerType,
+)
+from thunder.dynamo.repro_script_template import (
+    benchmark_multi_exe_code_template,
+    repro_code_template,
+    bsym_torch_compile_repro_template,
+)
+from thunder import last_traces, last_backward_traces
 
 
 if TYPE_CHECKING:
@@ -22,6 +34,9 @@
     from collections.abc import Sequence
 
     from thunder.dynamo.utils import ExampleInputMetaData
+    from thunder.dev_utils.utils import BenchmarkComparisonData
+    from thunder.core.trace import TraceCtx
+    from thunder.core.symbol import BoundSymbol
 
 
 def run_backward(fn, *args, **kwargs):
@@ -396,7 +411,7 @@ def __init__(self, graphs: list[torch.fx.GraphModule], graph_names: list[str] =
             graph_names = [f"graph{idx}" for idx in range(len(graphs))]
         self.fx_graph_reports: list[FXGraphReport] = [FXGraphReport(g, name) for name, g in zip(graph_names, graphs)]
 
-    def __str__(self):
+    def __repr__(self):
         return f"<FXReport with {len(self.fx_graph_reports)} FXGraphReports accessible via .fx_graph_reports>"
 
 
@@ -464,3 +479,313 @@ def helper_backend(gm, example_inputs):
     compiled(*args, **kwargs)
 
     return FXReport(graphs)
+
+
+class ThunderSplitGraphReport(FXGraphReport):
+    """
+    A report class representing a Thunder-split FX subgraph, extending :class:FXGraphReport.
+
+    This class encapsulates details about a subgraph that was split from the original
+    Dynamo FX graph due to Thunder-specific transformations.
+
+    Attributes:
+        graph: The Thunder-split FX graph.
+        graph_name: The name of the Thunder-split FX graph.
+        compiled_fn: The Thunder-compiled function corresponding to :attr:`graph`.
+        example_input: An example input used for execution.
+        thunder_options: Configuration options specific to :func:`thunder.jit`.
+        split_reason: The reason why :func:`thunder.dynamo.splitter._splitter` split the original graph.
+        fusion_reports (list[ThunderFusionReport]): A list of fusion reports for the
+            nvFusion regions generated when using the nvFuser executor.
+            See :class:`ThunderFusionReport` for more details.
+        fwd_trc: The forward trace, available only after calling :meth:`_create_thunder_traces`.
+        bwd_trc: The backward trace, available only after calling :meth:`_create_thunder_traces`.
+
+    For an example, see the documentation of :func:`analyze_thunder_splits`.
+    """
+
+    def __init__(
+        self,
+        graph: torch.fx.GraphModule,
+        graph_name: str,
+        compiled_fn: Callable,
+        example_input,
+        thunder_options: dict,
+        split_reason: str,
+    ):
+        super().__init__(graph, graph_name)
+        self.compiled_fn = compiled_fn
+        self.example_input = example_input
+        self.thunder_options = thunder_options
+        self.split_reason = split_reason
+
+        self.fusion_reports: list[ThunderFusionReport] = []
+        self.fwd_trc: TraceCtx = None
+        self.bwd_trc: TraceCtx = None
+
+    def __repr__(self):
+        return f"<ThunderSplitGraphReport with {len(self.fusion_reports)} ThunderFusionReport accessible via .fusion_reports>"
+
+    def _create_thunder_traces(self):
+        example_inputs = eval(f"[\n{chr(10).join(arg_like(a) for a in self.example_input)}]")
+        # Executes to get the trace
+        run_backward(self.compiled_fn, *example_inputs)
+        self.fwd_trc = last_traces(self.compiled_fn)[-1]
+        self.bwd_trc = last_backward_traces(self.compiled_fn)[-1]
+
+    def create_fusion_reports(self):
+        """
+        Runs the Thunder-compiled function to obtain the nvFusion definition
+        and generate the :class:`ThunderFusionReport` instance based on it.
+        """
+        self._create_thunder_traces()
+        for trace, prefix in [(self.fwd_trc, "forward"), (self.bwd_trc, "backward")]:
+            for bsym in trace.bound_symbols:
+                if bsym.sym.is_fusion and "nvFusion" in bsym.sym.name:
+                    self.fusion_reports.append(ThunderFusionReport(bsym, f"{self.graph_name}_{bsym.sym.name}_{prefix}"))
+
+    def write_thunder_repro(self, folder, use_benchmark=False, serialize_inputs=False):
+        thunder_ex_str = (
+            f"partial(thunder.jit, {self.thunder_options})" if self.thunder_options is None else "thunder.jit"
+        )
+        has_cuda_args = any(hasattr(arg, "device") and arg.device.type == "cuda" for arg in self.example_input)
+        import_str = ["import thunder", "from functools import partial"]
+        if has_cuda_args:
+            # Since Thunder compile options don't clearly indicate required imports,
+            # we include commonly used transforms by default.
+            import_str.extend(
+                [
+                    "from thunder.transforms.cudagraph import CUDAGraphTransform",
+                    "from thunder.dev_utils.nvtx_profile_transform import NvtxProfileTransform",
+                ]
+            )
+        if not use_benchmark:
+            super().write_repro(
+                folder,
+                f"{self.graph_name}_repro_thunder.py",
+                executor_str=thunder_ex_str,
+                import_str=import_str,
+                serialize_inputs=serialize_inputs,
+                inputs=self.example_input,
+                extra_comment_str=self.split_reason,
+            )
+            return
+
+        executor_names_list = ["thunder"]
+        executors = [thunder_ex_str]
+
+        super().write_benchmark_repro(
+            folder,
+            f"{self.graph_name}_benchmark_thunder.py",
+            executor_names_list,
+            executor_str=executors,
+            import_str=import_str,
+            serialize_inputs=serialize_inputs,
+            inputs=self.example_input,
+            extra_comment_str=self.split_reason,
+        )
+
+
+class ThunderFusionReport:
+    """
+    A report class representing a Thunder nvFusion region.
+
+    This class encapsulates information about a nvFusion region created during
+    Thunder execution, including the symbolic representation and its name.
+
+    Attributes:
+        nvfusion_bsym (BoundSymbol): The symbolic representation of the nvFusion region.
+        name (str): The name of the fusion region.
+
+    For an example, see the documentation of :func:`analyze_thunder_splits`.
+    """
+
+    def __init__(self, bsym: BoundSymbol, name: str):
+        self.nvfusion_bsym = bsym
+        self.name = name
+
+    def __repr__(self):
+        return f"<ThunderFusionReport of bound symbol\n{self.nvfusion_bsym}>"
+
+    def write_nvfuser_repro(self, folder, file_name=None):
+        folder = Path(folder)
+        folder.mkdir(exist_ok=True, parents=True)
+        nvfuser_callable = self.nvfusion_bsym._call_ctx[self.nvfusion_bsym.sym.name]
+        fd = nvfuser_callable.last_used
+
+        # The API for nvFuser version >=2.14
+        get_repro = getattr(fd, "repro_script_for", None)
+        # The legacy nvFuser API
+        if get_repro is None:
+            get_repro = getattr(fd, "getReproString", None)
+        if get_repro is None:
+            raise RuntimeError("The installed version of nvFuser does not support repro generation unless on crash.")
+
+        inputs = self.get_inputs()
+        repro_code = get_repro(inputs)
+        comment_str = f'"""\n{self.nvfusion_bsym}\n"""'
+
+        if file_name == None:
+            file_name = f"{self.name}_repro_nvfuser.py"
+        with open(folder / file_name, "w") as f:
+            print(comment_str, file=f)
+            print(repro_code, file=f)
+
+    def get_inputs(self):
+        return self.nvfusion_bsym._call_ctx[self.nvfusion_bsym.sym.name].last_inputs
+
+    def write_inductor_repro(self, folder: PathLike, file_name=None):
+        folder = Path(folder)
+        folder.mkdir(exist_ok=True, parents=True)
+        python_func = create_python_callable_from_bsym(self.nvfusion_bsym)
+        nvfusion_name = self.nvfusion_bsym.sym.name
+
+        inputs = self.get_inputs()
+        inputs = "[" + "".join(arg_like(inp) for inp in inputs) + "]"
+        program = bsym_torch_compile_repro_template.format(
+            python_func=python_func, func_name=nvfusion_name, inputs=inputs
+        )
+        if file_name == None:
+            file_name = f"{self.name}_repro_inductor.py"
+        with open(folder / file_name, "w") as f:
+            f.write(program)
+
+    def run_benchmark(self) -> BenchmarkComparisonData:
+        from thunder.dev_utils.utils import _benchmark_fusion_region_with_nvfuser_and_torch_compile
+
+        return _benchmark_fusion_region_with_nvfuser_and_torch_compile(self.nvfusion_bsym)
+
+
+class ThunderFXGraphReport(FXGraphReport):
+    """
+    A Thunder-specific report class for a Dynamo FX Graph, extending FXGraphReport,
+    providing the ability to save reproduction/benchmark scripts for the original FX graph.
+    Additionally, it includes information about Thunder-split subgraphs in `subgraph_reports`.
+
+    Attributes:
+        graph (torch.fx.GraphModule): The original Dynamo FX graph before being split by Thunder.
+        graph_name (str): The name of the original Dynamo FX graph.
+        split_reason (str): Reasons explaining why the subgraph was split.
+        subgraph_reports (list[ThunderSplitGraphReport]): A list of reports for each
+            Thunder-split FX graph. For more details, see :class:`ThunderSplitGraphReport`.
+
+    For an example, see the documentation of :func:`analyze_thunder_splits`.
+    """
+
+    def __init__(
+        self, gm: torch.fx.GraphModule, gm_name: str, split_reason: str, subgraph_reports: list[ThunderSplitGraphReport]
+    ):
+        super().__init__(gm, gm_name)
+        self.split_reason = split_reason
+        self.subgraph_reports: list[ThunderSplitGraphReport] = subgraph_reports
+
+    def __repr__(self):
+        return f"<ThunderFXGraphReport with {len(self.subgraph_reports)} ThunderSplitGraphReport accessible via .subgraph_reports>"
+
+
+def analyze_thunder_splits(
+    report: FXGraphReport,
+    **thunder_options,
+) -> ThunderFXGraphReport:
+    """
+    Generates a :class:`ThunderFXGraphReport` based on an :class:`FXGraphReport`.
+
+    The :class:`ThunderFXGraphReport` provides additional details about Thunder-specific splits
+    and nvFusion regions. For more details, see :class:`ThunderFXGraphReport`.
+
+    Example:
+        .. code-block:: python
+
+        import tempfile
+        from pathlib import Path
+        import torch
+        from thunder.dynamo.report import (
+            fx_report, FXReport, ThunderFXGraphReport, FXGraphReport,
+            ThunderSplitGraphReport, ThunderFusionReport, analyze_thunder_splits
+        )
+        from pathlib import Path
+
+        x = torch.ones(2, 2, device="cuda", requires_grad=True)
+
+        # Dynamo segments `foo` into two graphs. Each graph contains one Thunder-split graph,
+        # and each Thunder-split graph has one nvFusion region.
+        def foo(x):
+            x = x.exp()
+            torch._dynamo.graph_break()
+            y = torch.sinc(x) + torch.cos(x)
+            return y + 1
+
+        # If using `torch.compile` alone, you can stop here and query the reports in `FXReport`.
+        # For more details, see the example in :func:`fx_report`.
+        results: FXReport = fx_report(foo, x)
+
+        fx_graph_report: FXGraphReport
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = Path(tmpdir)
+            for idx, fx_graph_report in enumerate(results.fx_graph_reports):
+                # `ThunderFXGraphReport` extends `FXGraphReport`, providing the ability to save
+                # reproduction/benchmark scripts for the original FX graph. Additionally, it
+                # includes information about Thunder-split subgraphs in `subgraph_reports`.
+                thunder_fx_graph_report: ThunderFXGraphReport = analyze_thunder_splits(fx_graph_report)
+                # Saves a reproduction script for the original FX graph
+                thunder_fx_graph_report.write_thunder_repro(tmp_path)
+
+                thunder_split_report: ThunderSplitGraphReport
+                for thunder_split_report in thunder_fx_graph_report.subgraph_reports:
+                    split_folder = tmp_path / str(idx)
+                    thunder_split_report.write_eager_repro(split_folder)
+                    thunder_split_report.write_thunder_repro(split_folder)
+                    thunder_split_report.write_inductor_repro(split_folder)
+
+                    # If you are only interested in the Thunder-split FX graph, you can stop here.
+                    # If you want to inspect Thunder traces and nvFusion regions further, explicitly call
+                    # `ThunderSplitGraphReport.create_fusion_reports()` and analyze as shown below.
+                    thunder_split_report.create_fusion_reports()
+                    print(f"fwd_trace:\n{thunder_split_report.fwd_trc}\n")
+                    print(f"bwd_trace:\n{thunder_split_report.bwd_trc}\n")
+                    nvfusion_report: ThunderFusionReport
+                    for nvfusion_report in thunder_split_report.fusion_reports:
+                        nvfusion_report.write_nvfuser_repro(split_folder / "nvfusion")
+                        nvfusion_report.write_inductor_repro(split_folder / "nvfusion")
+                        bench_data = nvfusion_report.run_benchmark()
+                        print(bench_data)
+                        print("---"*10)
+
+    """
+    from thunder.dynamo.utils import remove_empty_autocast, recompile_graph, get_thunder_module_names
+    from thunder.dynamo.splitter import _splitter
+    from thunder import jit
+
+    # Splits the FX graph module using Thunder splitter
+    gm = remove_empty_autocast(report.graph)
+    # Dynamo uses lazy generation of the underlying Python code, so we need to
+    # force recompilation of the GraphModule before passing it to Thunder.
+    recompile_graph(gm)
+    _thunder_jit = partial(jit, **thunder_options, nv_store_fusion_inputs=True)
+    _, subgraph_info = _splitter(gm, _thunder_jit, torch.compile, _unused_sample_args=None)
+
+    thunder_module_names = [f"{report.graph_name}_{name}" for name in get_thunder_module_names(subgraph_info)]
+    original_modules_to_thunder_modules = (
+        [m, compiled_m]
+        for m, compiled_m in subgraph_info.submodule_to_compiled_functions.items()
+        if compiled_m.compiler == CompilerType.THUNDER
+    )
+    example_inputs = subgraph_info.thunder_compiled_fns_example_inputs
+    split_reason = get_split_reasons_string(subgraph_info)
+
+    subgraph_reports = []
+    for name, sub_gm_pair, example_input in zip(
+        thunder_module_names, original_modules_to_thunder_modules, example_inputs
+    ):
+        subgraph_reports.append(
+            ThunderSplitGraphReport(
+                sub_gm_pair[0],
+                name,
+                sub_gm_pair[1].compiled_fn,
+                example_input,
+                thunder_options,
+                split_reason,
+            )
+        )
+    report = ThunderFXGraphReport(report.graph, report.graph_name, split_reason, subgraph_reports)
+    return report
diff --git a/thunder/dynamo/repro_script_template.py b/thunder/dynamo/repro_script_template.py
index a951149f27..d0c0cee9f4 100644
--- a/thunder/dynamo/repro_script_template.py
+++ b/thunder/dynamo/repro_script_template.py
@@ -113,3 +113,28 @@ def test_{graph_name}(benchmark, executor, compute_type):
         compiled_model = executor(model)
     {call_benchmark}
 '''
+
+
+bsym_torch_compile_repro_template = """
+{python_func}
+
+from thunder.executors.torch_compile import make_compiled as make_torch_compile_callable
+import thunder.examine
+
+inputs = {inputs}
+
+jfn = thunder.jit({func_name})
+jfn(*inputs)
+
+trc = thunder.last_traces(jfn)[-1]
+fusion_symbols = thunder.examine.get_fusion_symbols(trc)
+assert len(fusion_symbols) == 1
+bsym = fusion_symbols[0]
+
+# NOTE: The nvFusion function cannot be compiled directly using `torch.compile`.
+# It must first be processed by Thunder into BoundSymbols and compiled with `make_torch_compile_callable`.
+# Additionally, it's recommended to visually verify that `bsym` matches the
+# `nvFusion` function above by printing it using `print(bsym)`.
+torch_compiled_callable = make_torch_compile_callable(bsym.subsymbols, bsym.flat_args, bsym.flat_outs)
+out = torch_compiled_callable(*inputs)
+"""
diff --git a/thunder/dynamo/utils.py b/thunder/dynamo/utils.py
index 9a0e779da8..8754bca7c4 100644
--- a/thunder/dynamo/utils.py
+++ b/thunder/dynamo/utils.py
@@ -729,3 +729,12 @@ def get_split_reasons_string(subgraph_info: SubgraphInfo) -> str:
     else:
         split_reason_str += "The original graph is not split, and is entirely run by Thunder.\n"
     return split_reason_str
+
+
+def get_thunder_module_names(subgraph_info: SubgraphInfo) -> list[str]:
+    thunder_module_names = []
+    for node in subgraph_info.split_graph_module.graph.nodes:
+        target = node.target
+        if isinstance(target, str) and target.startswith("thunder_"):
+            thunder_module_names.append(target)
+    return thunder_module_names


Test patch:
diff --git a/thunder/tests/test_dynamo.py b/thunder/tests/test_dynamo.py
index cd3f0b5462..9df96dafc5 100644
--- a/thunder/tests/test_dynamo.py
+++ b/thunder/tests/test_dynamo.py
@@ -1072,7 +1072,6 @@ def foo(x):
         assert "Failed to save reproducer" in captured.out
 
 
-@pytest.mark.parametrize("use_benchmark", (True, False), ids=("benchmark", "repro"))
 @instantiate(
     dtypes=NOTHING,
     executors=[DynamoThunderExecutor],
@@ -1123,7 +1122,7 @@ def check(file_name, cmd):
     assert len(py_files) == 4
 
     for file in py_files:
-        check(tmp_path / file, cmd)
+        check(file, cmd)
 
 
 def test_leak_on_unsupported_thunder_operator():
@@ -1151,3 +1150,45 @@ def call_thunderfx_on_leaking_fn():
 
     t_weak_ref = call_thunderfx_on_leaking_fn()
     assert t_weak_ref() is None
+
+
+@requiresCUDA
+def test_thunder_specific_reports(tmp_path):
+    from thunder.dynamo.report import fx_report, analyze_thunder_splits
+
+    x = torch.ones(2, 2, device="cuda", requires_grad=True)
+
+    def foo(x):
+        # torch.sinc has automatic fallback registered,
+        # so that operation will be given to inductor.
+        x = x.exp()
+        torch._dynamo.graph_break()
+        y = torch.sinc(x) + torch.cos(x)
+        return y + 1
+
+    results = fx_report(foo, x)
+    for idx, fx_graph_report in enumerate(results.fx_graph_reports):
+        thunder_fx_graph_report = analyze_thunder_splits(fx_graph_report)
+        thunder_fx_graph_report.write_thunder_repro(tmp_path)
+        for thunder_split_report in thunder_fx_graph_report.subgraph_reports:
+            split_folder = tmp_path / str(idx)
+            thunder_split_report.write_eager_repro(split_folder)
+            thunder_split_report.write_thunder_repro(split_folder)
+            thunder_split_report.write_inductor_repro(split_folder)
+            thunder_split_report.create_fusion_reports()
+            for nvf in thunder_split_report.fusion_reports:
+                nvf.write_nvfuser_repro(split_folder / "nvfusion")
+                nvf.write_inductor_repro(split_folder / "nvfusion")
+                bench_data = nvf.run_benchmark()
+
+    def check(file_name, cmd):
+        cmd = cmd + [file_name]
+        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+        assert result.returncode == 0, f"Script {file_name} failed: {result}"
+
+    cmd = [sys.executable]
+    py_files = list(tmp_path.rglob("*.py"))
+    assert len(py_files) == 16
+
+    for file in py_files:
+        check(file, cmd)