From b857701d540f8bb964311e0bad6526374de2c965 Mon Sep 17 00:00:00 2001
From: Masaki Kozuki <mkozuki@nvidia.com>
Date: Fri, 27 Sep 2024 00:54:56 -0700
Subject: [PATCH] try removing `ltorch.copy_`

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
---
 thunder/executors/nvfuserex_impl.py |  3 +--
 thunder/tests/test_inplace_copy.py  | 26 --------------------------
 thunder/torch/__init__.py           |  5 -----
 3 files changed, 1 insertion(+), 33 deletions(-)

diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py
index a13691103f..4ae4b4f2ed 100644
--- a/thunder/executors/nvfuserex_impl.py
+++ b/thunder/executors/nvfuserex_impl.py
@@ -2071,8 +2071,7 @@ def copy_(
 ) -> Any:
     nvcopy_from = getnv(copy_from, fd, lc_to_nv_map)
     nvcopy_to = getnv(copy_to, fd, lc_to_nv_map)
-    alias_output = fd.ops.set(nvcopy_from)
-    fd.add_output(alias_output, alias_input=nvcopy_to)
+    fd.add_output(nvcopy_from, alias_input=nvcopy_to)
     return nvcopy_to
 
 
diff --git a/thunder/tests/test_inplace_copy.py b/thunder/tests/test_inplace_copy.py
index 952df2faf0..feb49fe77a 100644
--- a/thunder/tests/test_inplace_copy.py
+++ b/thunder/tests/test_inplace_copy.py
@@ -152,29 +152,3 @@ def func3(x, y):
             match=r"If you are sure you don't want to use this check, it can be disabled by setting `disable_inplace_copy_check=True` in `thunder.jit`.$",
         ):
             traced_foo(a, b)
-
-
-@instantiate(executors=(nvFuserExecutor,), dtypes=(thunder.float32,))
-def test_inplace_copy_dst_copy_returned_issue_1109(executor, device, dtype):
-    def func(T0):
-        T1 = torch.sin(T0)
-        T0.copy_(T1)  # destination.copy_(source)
-        T2 = torch.cos(T1)
-        T0.copy_(T2)
-        # T1 & T2 should be returned as separate buffer, instead of sharing
-        # storage with T0
-        return T1, T2
-
-    tdtype = ttorch.to_torch_dtype(dtype)
-    # This pattern is unsafe in general. Disabling sanity check to silence
-    # exception for testing
-    traced_foo = executor.make_callable(func, disable_inplace_copy_check=True)
-    a = make_tensor((4, 4), device=device, dtype=tdtype)
-    a_ref = a.clone()
-
-    o_thunder = traced_foo(a)
-    o_eager = func(a_ref)
-
-    assert_close(a_ref, a)
-    for o, o_ref in zip(o_thunder, o_eager):
-        assert_close(o, o_ref)
diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py
index 3188855323..18ba597973 100644
--- a/thunder/torch/__init__.py
+++ b/thunder/torch/__init__.py
@@ -1945,11 +1945,6 @@ def copysign_(a, b, /):
     return prims.copy_(copysign(a, b), a)
 
 
-@torchsymbol(torch.Tensor.copy_, is_method=True)  # , tags=(prims.OpTags.IN_PLACE,))
-def copy_(a, b, /):
-    return prims.copy_(b, a)
-
-
 # TODO Implement div
 @torchsymbol(torch.div, is_method=True)
 def div(