Add tests for use_dtensor_execute flag

Co-authored-by: IvanYashchuk <19621411+IvanYashchuk@users.noreply.github.com>

Author: Copilot

thunder/tests/distributed/test_dtensor.py

@@ -459,6 +459,47 @@ def test_dtensor_opinfo(self, op: OpInfo, executor):
         assert tested_sample_count > 0, f"test_dtensor_opinfo:No samples tested for {op.name} with {executor} executor"
 
 
+    @common_utils.parametrize("executor", tuple(executors_map.keys()))
+    def test_use_dtensor_execute_flag_with_dtensor(self, executor):
+        """Test that use_dtensor_execute flag is set to True for DTensor inputs."""
+        import thunder.examine as examine
+
+        num_devices = self.world_size
+        mesh = DeviceMesh("cuda", list(range(num_devices)))
+        dim_size = 16
+
+        in_dtensor = distribute_tensor(torch.randn(dim_size, dim_size, requires_grad=True), mesh, [Shard(0)])
+        w_dtensor = distribute_tensor(torch.randn(dim_size, dim_size, requires_grad=True), mesh, [Shard(0)])
+
+        def fn(x, w):
+            return torch.mul(x, w)
+
+        tmodel = thunder.jit(fn, executors=executors_map[executor].executors_list())
+        result = tmodel(in_dtensor, w_dtensor)
+
+        traces = thunder.last_traces(tmodel)
+        extrace = traces[-1]
+        fusions = examine.get_fusion_symbols(extrace)
+
+        # Verify there is at least one fusion
+        assert len(fusions) > 0
+
+        # Check that use_dtensor_execute is True for DTensor inputs
+        for fusion_bsym in fusions:
+            _, call_ctx, _ = fusion_bsym.gather_ctxs()
+            fusion_name = fusion_bsym.sym.name
+            fdw = call_ctx.get(fusion_name)
+            assert fdw is not None, f"FusionDefinitionWrapper not found for {fusion_name}"
+            assert hasattr(
+                fdw, "use_dtensor_execute"
+            ), "FusionDefinitionWrapper should have use_dtensor_execute attribute"
+            assert fdw.use_dtensor_execute is True, "use_dtensor_execute should be True for DTensor inputs"
+
+        # Verify the result is correct
+        expected = torch.mul(in_dtensor, w_dtensor)
+        torch.testing.assert_close(result, expected)
+
+
 common_utils.instantiate_parametrized_tests(DTensorTest)
 
 if __name__ == "__main__":

thunder/tests/test_nvfuser.py

@@ -1040,3 +1040,40 @@ def test_scatter(executor, device: str, dtype: dtypes.dtype):
         assert len(fusion_bsyms) == 1
         outside_fusion_syms = ["unpack_trivial", "python_return"]
         assert {el.sym.name for el in fw_trace.bound_symbols if not el.sym.is_fusion} == set(outside_fusion_syms)
+
+
+@instantiate(
+    executors=(nvFuserExecutor,),
+    dtypes=(thunder.float32,),
+)
+def test_use_dtensor_execute_flag(executor, device: str, dtype: dtypes.dtype):
+    """Test that use_dtensor_execute flag is set correctly at trace construction time."""
+    torch_dtype = ltorch.to_torch_dtype(dtype)
+    a = make_tensor((2, 2), device=device, dtype=torch_dtype)
+    b = make_tensor((2, 2), device=device, dtype=torch_dtype)
+
+    def foo(a, b):
+        return a + b
+
+    cfoo = thunder.jit(foo)
+    result = cfoo(a, b)
+
+    traces = thunder.last_traces(cfoo)
+    extrace = traces[-1]
+    fusions = examine.get_fusion_symbols(extrace)
+
+    # Verify there is at least one fusion
+    assert len(fusions) > 0
+
+    # Check that use_dtensor_execute is False for regular tensors
+    for fusion_bsym in fusions:
+        _, call_ctx, _ = fusion_bsym.gather_ctxs()
+        fusion_name = fusion_bsym.sym.name
+        fdw = call_ctx.get(fusion_name)
+        assert fdw is not None, f"FusionDefinitionWrapper not found for {fusion_name}"
+        assert hasattr(fdw, "use_dtensor_execute"), "FusionDefinitionWrapper should have use_dtensor_execute attribute"
+        assert fdw.use_dtensor_execute is False, "use_dtensor_execute should be False for regular tensors"
+
+    # Verify the result is correct
+    expected = a + b
+    torch.testing.assert_close(result, expected)