From fca5f612f25fd71b7bfec4009e42224aa225884f Mon Sep 17 00:00:00 2001
From: PanZezhong <panzezhong@qiyuanlab.com>
Date: Tue, 25 Aug 2026 01:33:14 +0000
Subject: [PATCH] feat: support cross-device copy_from

---
 python/infinicore/tensor.py   |  5 ++++-
 src/infinicore/tensor/copy.cc | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py
index bbe801f93..d10929dce 100644
--- a/python/infinicore/tensor.py
+++ b/python/infinicore/tensor.py
@@ -190,7 +190,10 @@ def strided_from_blob(data_ptr, size, strides, *, dtype=None, device=None):
 
 def from_torch(torch_tensor) -> Tensor:
     infini_type = to_infinicore_dtype(torch_tensor.dtype)
-    infini_device = infinicore.device(torch_tensor.device.type, 0)
+    device_index = torch_tensor.device.index
+    infini_device = infinicore.device(
+        torch_tensor.device.type, 0 if device_index is None else device_index
+    )
     return Tensor(
         _infinicore.from_blob(
             torch_tensor.data_ptr(),
diff --git a/src/infinicore/tensor/copy.cc b/src/infinicore/tensor/copy.cc
index 1297d9f8c..02151f5e5 100644
--- a/src/infinicore/tensor/copy.cc
+++ b/src/infinicore/tensor/copy.cc
@@ -23,9 +23,11 @@ void TensorImpl::copy_from(Tensor src) {
             "Cannot copy from tensor with different shape. Src: " + src->info() + " Dst: " + this->info());
     }
     if (this->device() == src->device()) {
+        context::setDevice(this->device());
         op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), src);
     } else {
         if (!src->is_contiguous()) {
+            context::setDevice(src->device());
             src = src->contiguous();
         }
 
@@ -50,6 +52,20 @@ void TensorImpl::copy_from(Tensor src) {
                 context::memcpyH2D(local_src->data(), src->data(), copy_size);
                 op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), local_src);
             }
+        } else {
+            if (this->device().getType() != src->device().getType()) {
+                throw std::runtime_error(
+                    "Cannot copy directly between different accelerator types. Src: " + src->info() +
+                    " Dst: " + this->info());
+            }
+            context::setDevice(this->device());
+            if (this->is_contiguous()) {
+                context::memcpyD2D(this->data(), src->data(), copy_size);
+            } else {
+                auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device());
+                context::memcpyD2D(local_src->data(), src->data(), copy_size);
+                op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), local_src);
+            }
         }
     }
 }