Add prims.clone (#86705)

IvanYashchuk · pytorchmergebot · commit cd7c86eaa468 · 2022-10-12T18:22:00.000Z
This simple PR adds `clone` as a primitive. Current implementation of `clone` is not supported with nvFuser executor because of `empty_like` + `copy_to`. Pull Request resolved: #86705 Approved by: https://github.com/mruberry
diff --git a/torch/_prims/__init__.py b/torch/_prims/__init__.py
@@ -159,6 +159,7 @@
     #
     # Data conversion and movement prims
     #
+    "clone",
     "convert_element_type",
     "device_put",
     "item",
@@ -595,6 +596,40 @@ def _conj_physical_meta(input: TensorLikeType) -> TensorLikeType:
     return_type=RETURN_TYPE.NEW,
 )
 
+
+def _clone_meta(
+    input: TensorLikeType, *, memory_format: torch.memory_format = torch.preserve_format
+) -> TensorLikeType:
+    if memory_format != torch.preserve_format:
+        return torch.empty(
+            input.shape,
+            dtype=input.dtype,
+            layout=input.layout,
+            device=input.device,
+            requires_grad=input.requires_grad,
+            memory_format=memory_format,
+        )
+
+    # memory_format == torch.preserve_format
+    strides = utils.compute_elementwise_output_strides(input)
+    return torch.empty_strided(
+        input.shape,
+        strides,
+        dtype=input.dtype,
+        layout=input.layout,
+        device=input.device,
+        requires_grad=input.requires_grad,
+    )
+
+
+clone = _make_prim(
+    schema="clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor",
+    meta=_clone_meta,
+    impl_aten=torch.clone,
+    doc="Returns the copy of a tensor",
+    return_type=RETURN_TYPE.NEW,
+)
+
 digamma = _make_elementwise_unary_prim(
     "digamma",
     impl_aten=torch.digamma,
diff --git a/torch/_prims/nvfuser_prims.py b/torch/_prims/nvfuser_prims.py
@@ -42,6 +42,7 @@
     "atanh",
     "cos",
     "cosh",
+    "clone",
     "bitwise_not",
     "ceil",
     "erf",
@@ -322,9 +323,14 @@ def _amin_nvfuser(
     return fd.ops.min(a, dims, keep_dims)
 
 
+def _clone_nvfuser(fd: Any, input: TensorLikeType, *, memory_format=None):
+    return fd.ops.set(input)
+
+
 _nvfuser_impls["native_batch_norm"] = _native_batch_norm_nvfuser
 _nvfuser_impls["broadcast_in_dim"] = _broadcast_in_dim_nvfuser
 _nvfuser_impls["convert_element_type"] = _convert_element_type_nvfuser
+_nvfuser_impls["clone"] = _clone_nvfuser
 _nvfuser_impls["transpose"] = _transpose_nvfuser
 _nvfuser_impls["squeeze"] = _squeeze_nvfuser
 _nvfuser_impls["view_of"] = _view_of_nvfuser
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
@@ -1675,10 +1675,7 @@ def where(
 def clone(
     a: TensorLikeType, *, memory_format: torch.memory_format = torch.preserve_format
 ) -> TensorLikeType:
-    result = torch.empty_like(
-        a, requires_grad=a.requires_grad, memory_format=memory_format
-    )
-    copy_to(result, a)
+    result = prims.clone(a, memory_format=memory_format)
     return result
 
 
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -17578,7 +17578,6 @@ def reference_flatten(input, start_dim=0, end_dim=-1):
     PythonRefInfo(
         "_refs.clone",
         torch_opinfo_name="clone",
-        supports_nvfuser=False,
     ),
     #
     # View & Shape OpInfos

Original file line number	Diff line number	Diff line change
`@@ -17578,7 +17578,6 @@ def reference_flatten(input, start_dim=0, end_dim=-1):`
`17578`	`17578`	`PythonRefInfo(`
`17579`	`17579`	`"_refs.clone",`
`17580`	`17580`	`torch_opinfo_name="clone",`
`17581`		`- supports_nvfuser=False,`
`17582`	`17581`	`),`
`17583`	`17582`	`#`
`17584`	`17583`	`# View & Shape OpInfos`