Merge branch 'main' into do-not-return-deepcopy-of-GraphModule

Author: Masato Shinokawa

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

thunder/tests/opinfos.py

@@ -2192,6 +2192,14 @@ def imag_error_generator(op, device, **kwargs):
 elementwise_unary_ops.append(clone_opinfo)
 
 
+square_opinfo = OpInfo(
+    ltorch.square,
+    sample_input_generator=elementwise_unary_generator,
+    torch_reference=_elementwise_unary_torch(torch.square),
+)
+elementwise_unary_ops.append(square_opinfo)
+
+
 # Puts all opinfos into the "opinfos" list
 opinfos.extend(elementwise_unary_ops)
 

thunder/torch/__init__.py

@@ -2248,6 +2248,25 @@ def threshold_(a: TensorProxy, /, threshold: float, value: float) -> TensorLike:
 _inplace_to_out_of_place[threshold_] = threshold, -1
 
 
+@torchsymbol(torch.square, is_method=True)
+def square(a):
+    if isinstance(dtypes.to_dtype(a), dtypes.bool_):
+        a = clang.maybe_convert_to_dtype(a, dtypes.int64)
+    return a * a
+
+
+@torchsymbol(torch.square_, is_method=True, tags=(prims.OpTags.IN_PLACE,))
+def square_(a):
+    utils.check(
+        not isinstance(dtypes.to_dtype(a), dtypes.bool_),
+        lambda: f"Result type of {dtypes.int64} cannot be stored into {dtypes.to_dtype(a)}",
+    )
+    return _copy_(a, square(a))
+
+
+_inplace_to_out_of_place[square_] = square, -1
+
+
 #
 # Elementwise binary operations
 #

thunder/torch/default_torch_ops.py

@@ -252,7 +252,6 @@
         torch.split_with_sizes,
         torch.split_with_sizes_copy,
         torch.spmm,
-        torch.square,
         torch.squeeze_copy,
         torch.sspaddmm,
         torch.std_mean,
@@ -548,7 +547,6 @@
         torch.Tensor.sparse_dim,
         torch.Tensor.sparse_mask,
         torch.Tensor.split_with_sizes,
-        torch.Tensor.square,
         torch.Tensor.sspaddmm,
         torch.Tensor.stft,
         torch.Tensor.subtract,