Merge branch 'main' into te-v2-reset-transform

Author: riccardofelluga

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
@@ -1393,3 +1394,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/executors/nvfuserex_impl.py

@@ -71,6 +71,12 @@
 from thunder.extend import FUEL_LEVEL, FusionExecutor, register_executor
 from thunder.executors.nvfuserex import nvfuser_version
 
+
+DTENSOR_SUPPORTED_VERSION = LooseVersion("0.2.28")
+if nvfuser_version() >= DTENSOR_SUPPORTED_VERSION:
+    import nvfuser_direct as nvfd
+    from nvfuser_direct import FusionDefinition as DirectFusionDefinition
+
 # NOTE This impl file is here because nvFuser may not be available, so it's imported conditionally
 #   by nvfuserex.py when nvFuser is available.
 import nvfuser
@@ -241,41 +247,34 @@ def get_translator(bsym: BoundSymbol) -> Callable:
     return _translation_map[bsym.sym.id]
 
 
-class MultiDeviceFusionDefinition(FusionDefinition):
-    def __init__(self, define_fn: Callable[[FusionDefinition], None], in_dtensors: list[DTensorProxy], max_length: int):
-        super().__init__(max_length=max_length)
-        self._in_dtensors = in_dtensors
-        self._define_fn = define_fn
+def register_dtensor_supported(prim_id: int, fn: Callable, checker_fn: Callable) -> None:
+    if nvfuser_version() < DTENSOR_SUPPORTED_VERSION:
+        # Only register dtensor ops if supported version is available.
+        return
 
-    def definition(self) -> None:
-        self._define_fn(self)
+    register_supported(prim_id, fn, checker_fn)
 
-    def _find_tensor_by_index(self, index: int) -> nvfuser.Tensor:
-        for t in self.sched.tensors():
-            if t.index == index:
-                return t
-        return None
 
-    def multidevice_schedule(self) -> None:
-        for in_tensor_index, in_dtensor in zip(self.inputs(), self._in_dtensors):
-            in_tensor = self._find_tensor_by_index(in_tensor_index)
+def multidevice_schedule(fd: FusionDefinition, in_dtensors: list[Proxy]) -> None:
+    for in_tv, in_dtensor in zip(fd.fusion.inputs(), in_dtensors):
+        assert isinstance(in_dtensor, DTensorProxy)
+        # Set the device mesh.
+        assert in_dtensor.device_mesh.ndim == 1, "nvFuser's Python API only supports 1D meshes."
+        mesh = nvfd.multidevice.DeviceMesh(in_dtensor.device_mesh.mesh.tolist())
 
-            # Set the device mesh.
-            utils.check(in_dtensor.device_mesh.ndim == 1, lambda: "nvFuser's Python API only supports 1D meshes.")
-            mesh = nvfuser.DeviceMesh(in_dtensor.device_mesh.mesh.tolist())
+        in_tv.set_device_mesh(mesh)
 
-            self.sched._set_device_mesh(in_tensor, mesh)
+        assert len(in_dtensor.placements) == 1, "nvFuser's Python API only supports 1D meshes."
 
-            # Split and parallelize.
-            utils.check(len(in_dtensor.placements) == 1, lambda: "nvFuser's Python API only supports 1D meshes.")
-            # When the mesh is multi-dimensional, iterate through the
-            # placements in descending order of Placement.dim.
-            placement: Placement = in_dtensor.placements[0]
-            if placement.is_shard():
-                dim = cast(Shard, placement).dim
-                self.sched.split(in_tensor, dim, mesh.size, False)
-                self.sched.parallelize(in_tensor, dim, nvfuser.ParallelType.mesh_x)
-                self.sched.set_allocation_as_loop(in_tensor)
+        # Split and parallelize.
+        # When the mesh is multi-dimensional, iterate through the
+        # placements in descending order of Placement.dim.
+        placement: Placement = in_dtensor.placements[0]
+        if placement.is_shard():
+            dim = cast(Shard, placement).dim
+            in_tv.split(dim, mesh.size, inner_split=False)
+            in_tv.axis(dim).parallelize(nvfd.ParallelType.mesh_x)
+            in_tv.set_allocation_domain(in_tv.get_loop_domain(), new_contiguity=True)
 
 
 def create_fd(
@@ -376,10 +375,13 @@ def check_dtensor_tracing_and_runtime_metadata(inp):
             lambda: "nvfuser: Expected runtime and tracing metadata to be the same for DTensor.",
         )
 
-        fd = MultiDeviceFusionDefinition(definition, sorted_unique_inputs, max_length=MAX_LENGTH)
+        fd = DirectFusionDefinition()
         # Device may be set in one of the "factory" methods like full, iota, or uniform
         # NOTE: This should be called before defining because a factory method may look-up at `_selected_device` while being defined.
         fd._selected_device = None
+        with fd:
+            definition(fd)
+            multidevice_schedule(fd, sorted_unique_inputs)
     else:
         # NOTE nvFuser's default max length is 1024 operations at the time of this writing
         #   This arbitrarily increases it to 9999
@@ -535,28 +537,10 @@ def __call__(self, *args):
         if self.store_inputs:
             self.last_inputs = args
 
-        if hasattr(fd, "multidevice_schedule"):
+        if dist.is_available() and any(isinstance(t, torch.distributed.tensor.DTensor) for t in args):
             with annotate_for_profile(self.name):
-                in_tensors = [in_dtensor.to_local() for in_dtensor in args]
-                out_tensors, out_shardings = fd.execute(
-                    in_tensors,
-                    device=fd._selected_device,
-                    save_repro_inputs=self.save_fake_inputs,
-                    _enable_options=self.enable_options,
-                    _disable_options=self.disable_options,
-                )
-
-                assert len(out_tensors) == len(out_shardings)
-                out_dtensors: list[DTensor] = []
-                for out_tensor, out_sharding in zip(out_tensors, out_shardings):
-                    mesh = dist.device_mesh.init_device_mesh("cuda", (out_sharding.mesh.size,))
-                    placements: list[Placement] = []
-                    for parallel_type in [nvfuser.ParallelType.mesh_x]:
-                        axis: int = out_sharding.axis_sharded_on(parallel_type)
-                        placements.append(Replicate() if axis == -1 else Shard(axis))
-                    out_dtensors.append(DTensor.from_local(out_tensor, mesh, placements))
-
-                return out_dtensors
+                output = nvfd.execute_with_dtensors(fd, args)
+                return output
         else:
             with annotate_for_profile(self.name):
                 return fd.execute(
@@ -1917,7 +1901,7 @@ def mul(a: TensorProxy | Number, b: TensorProxy | Number, *, fd: FusionDefinitio
 
 
 register_supported(PrimIDs.MUL, mul, _elementwise_binary_check)
-register_supported(dtensor_mul_prim.id, mul, _elementwise_binary_check)
+register_dtensor_supported(dtensor_mul_prim.id, mul, _elementwise_binary_check)
 
 
 def ne(a: TensorProxy | Number, b: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:

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,