Add aten::clone and simple cases

* Add aten::clone and simple cases Signed-off-by: Feng Yuan <feng1.yuan@intel.com> * Fix Loops legacy code-path Signed-off-by: Feng Yuan <feng1.yuan@intel.com> --------- Signed-off-by: Feng Yuan <feng1.yuan@intel.com>
intel · Mar 8, 2024 · 9cb104a · 9cb104a
1 parent 6340065
commit 9cb104a
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 2 deletions.
diff --git a/src/aten/TensorFactories.cpp b/src/aten/TensorFactories.cpp
@@ -28,9 +28,14 @@ Tensor empty_strided_xpu(IntArrayRef size, IntArrayRef stride, c10::optional<Sca
   return result;
 }
 
+Tensor clone_xpu(const Tensor& self, c10::optional<MemoryFormat> memory_format) {
+  return at::native::clone(self, memory_format);
+}
+
 TORCH_LIBRARY_IMPL(aten, XPU, m) {
   m.impl(TORCH_SELECTIVE_NAME("aten::empty.memory_format"), TORCH_FN(at::native::empty_xpu));
   m.impl(TORCH_SELECTIVE_NAME("aten::empty_strided"), TORCH_FN(at::native::empty_strided_xpu));
+  m.impl(TORCH_SELECTIVE_NAME("aten::clone"), TORCH_FN(at::native::clone_xpu));
 }
 
 } // namespace at::native
diff --git a/src/aten/sycl/Loops.h b/src/aten/sycl/Loops.h
@@ -43,13 +43,13 @@ inline void elementwise_kernel_helper(func_t f, policy_t policy) {
 template <int vec_size, typename func_t>
 struct ElementwiseKernel {
   void operator()(sycl::nd_item<1> item) const {
-    int grpsz = item.get_local_range(0);
+    int glbsz = item.get_global_range(0);
     int gid = item.get_global_linear_id();
   #pragma unroll
     for (int i = 0; i < vec_size; i++) {
       if (gid < numel_) {
         f_(gid);
-        gid += grpsz;
+        gid += glbsz;
       }
     }
   };

diff --git a/test/python/examples/test_copy.py b/test/python/examples/test_copy.py
@@ -0,0 +1,21 @@
+import torch
+from torch.testing._internal.common_utils import TestCase
+
+cpu_device = torch.device("cpu")
+xpu_device = torch.device("xpu")
+
+
+class TestSimpleCopy(TestCase):
+    def test_copy_and_clone(self, dtype=torch.float):
+        a_cpu = torch.randn(16, 64, 28, 28)
+        b_cpu = torch.randn(16, 64, 28, 28)
+        a_xpu = a_cpu.to(xpu_device)
+        b_xpu = b_cpu.to(xpu_device)
+        # naive
+        b_cpu.copy_(a_cpu)
+        b_xpu.copy_(a_xpu)
+        self.assertEqual(b_cpu, b_xpu.to(cpu_device))
+        # clone + permutation
+        b_cpu = a_cpu.clone(memory_format=torch.channels_last)
+        b_xpu = a_xpu.clone(memory_format=torch.channels_last)
+        self.assertEqual(b_cpu, b_xpu.to(cpu_device))