Merge branch 'main' into remove_none_helper

Author: beverlylytle

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(
@@ -1906,6 +1890,9 @@ def le(a: TensorProxy | Number, b: TensorProxy | Number, *, fd: FusionDefinition
     return fd.ops.le(nva, nvb)
 
 
+register_supported(PrimIDs.LE, le, _elementwise_binary_check)
+
+
 def lt(a: TensorProxy | Number, b: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:
     nva = getnv(a, fd, lc_to_nv_map)
     nvb = getnv(b, fd, lc_to_nv_map)
@@ -1924,7 +1911,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,