Add update_inplace overload accepting OrtValue for device-to-device copy

onnxruntime/python/onnxruntime_inference_collection.py

@@ -13,10 +13,11 @@
 from enum import IntEnum
 from typing import Any
 
+import numpy as np
+
 from onnxruntime.capi import _pybind_state as C
 
 if typing.TYPE_CHECKING:
-    import numpy as np
     import numpy.typing as npt
 
     import onnxruntime
@@ -1212,8 +1213,6 @@ def __array__(self, dtype=None, copy=None) -> np.ndarray:
             If ``None`` (default), a copy will be made only if needed.
         :return: A numpy array with the same data as the OrtValue.
         """
-        import numpy as np  # noqa: PLC0415
-
         arr = self.numpy()
 
         if copy is not None:
@@ -1302,15 +1301,25 @@ def from_dlpack(cls, data, /) -> OrtValue:
 
         return cls(C.OrtValue.from_dlpack(capsule, is_bool))
 
-    def update_inplace(self, np_arr) -> None:
+    def update_inplace(self, data) -> None:
         """
-        Update the OrtValue in place with a new Numpy array. The numpy contents
-        are copied over to the device memory backing the OrtValue. It can be used
-        to update the input valuess for an InferenceSession with CUDA graph
-        enabled or other scenarios where the OrtValue needs to be updated while
-        the memory address can not be changed.
+        Update the OrtValue in place. The source data is copied over to the device
+        memory backing the OrtValue. It can be used to update the input values for
+        an InferenceSession with CUDA graph enabled or other scenarios where the
+        OrtValue needs to be updated while the memory address can not be changed.
+
+        :param data: The source data, which can be a Numpy array or another OrtValue.
+            When an OrtValue is provided, data can be copied between devices (e.g.,
+            GPU to GPU) without going through the CPU.
         """
-        self._ortvalue.update_inplace(np_arr)
+        if isinstance(data, OrtValue):
+            self._ortvalue.update_inplace(data._ortvalue)
+            return
+
+        if not isinstance(data, np.ndarray):
+            raise TypeError("data must be a numpy.ndarray or an OrtValue.")
+
+        self._ortvalue.update_inplace(data)
 
 
 def copy_tensors(src: Sequence[OrtValue], dst: Sequence[OrtValue], stream=None) -> None:

onnxruntime/python/onnxruntime_pybind_mlvalue.cc

@@ -1071,5 +1071,111 @@ void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const
   }
 }
 
+void UpdateOrtValueInplace(OrtValue& dst, const OrtValue& src) {
+  if (!dst.IsTensor()) {
+    throw std::runtime_error("Inplace update of OrtValues is only supported for Tensors");
+  }
+  if (!src.IsTensor()) {
+    throw std::runtime_error("The source OrtValue must contain a Tensor");
+  }
+
+  const auto& dst_tensor = dst.Get<Tensor>();
+  const auto& src_tensor = src.Get<Tensor>();
+
+  if (dst_tensor.DataType() != src_tensor.DataType()) {
+    throw std::runtime_error("The source and destination OrtValues must have the same data type");
+  }
+
+  if (dst_tensor.Shape().Size() != src_tensor.Shape().Size()) {
+    throw std::runtime_error("The source and destination OrtValues must have the same size");
+  }
+
+  if (dst_tensor.IsDataTypeString()) {
+    throw std::runtime_error("Inplace update of string tensors is not supported");
+  }
+
+  size_t bytes = 0;
+  auto status = Tensor::CalculateTensorStorageSize(dst_tensor.DataType(), dst_tensor.Shape(), 0, bytes);
+  if (!status.IsOK()) {
+    throw std::runtime_error(status.ErrorMessage());
+  }
+
+  const auto src_device = src_tensor.Location().device;
+  const auto dst_device = dst_tensor.Location().device;
+
+  void* dst_ptr = dst.GetMutable<Tensor>()->MutableDataRaw();
+  const void* src_ptr = src_tensor.DataRaw();
+
+  if (src_device.UsesCpuMemory() && dst_device.UsesCpuMemory()) {
+    memcpy(dst_ptr, src_ptr, bytes);
+  } else {
+    auto copy_fn = CreateDataTransferMemCpy(src_device, dst_device);
+    if (!copy_fn) {
+      // Fall back to built-in EP copy functions.
+      // Gate each path on (Type, VendorId) so that builds with multiple GPU EPs
+      // (e.g. CUDA + DML) route through the correct backend.
+      const auto is_cuda_device = [](const OrtDevice& device) {
+        return device.Type() == OrtDevice::GPU && device.Vendor() == OrtDevice::VendorIds::NVIDIA;
+      };
+      const auto is_migraphx_device = [](const OrtDevice& device) {
+        return device.Type() == OrtDevice::GPU && device.Vendor() == OrtDevice::VendorIds::AMD;
+      };
+      const auto is_dml_device = [](const OrtDevice& device) {
+        return device.Type() == OrtDevice::GPU && device.Vendor() == OrtDevice::VendorIds::MICROSOFT;
+      };
+      const auto is_cann_device = [](const OrtDevice& device) {
+        return device.Type() == OrtDevice::NPU && device.Vendor() == OrtDevice::VendorIds::HUAWEI;
+      };
+#ifdef USE_CUDA
+      if (is_cuda_device(src_device) && is_cuda_device(dst_device)) {
+        auto data_transfer = GetGPUDataTransfer();
+        ORT_THROW_IF_ERROR(data_transfer->CopyTensor(src_tensor, *dst.GetMutable<Tensor>()));
+        return;
+      }
+      if (src_device.UsesCpuMemory() && is_cuda_device(dst_device)) {
+        CpuToCudaMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+      if (is_cuda_device(src_device) && dst_device.UsesCpuMemory()) {
+        CudaToCpuMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+#endif
+#if USE_MIGRAPHX
+      if (src_device.UsesCpuMemory() && is_migraphx_device(dst_device)) {
+        CpuToMIGraphXMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+      if (is_migraphx_device(src_device) && dst_device.UsesCpuMemory()) {
+        MIGraphXToCpuMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+#endif
+#if USE_DML
+      if (src_device.UsesCpuMemory() && is_dml_device(dst_device)) {
+        CpuToDmlMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+      if (is_dml_device(src_device) && dst_device.UsesCpuMemory()) {
+        DmlToCpuMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+#endif
+#ifdef USE_CANN
+      if (src_device.UsesCpuMemory() && is_cann_device(dst_device)) {
+        CpuToCannMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+      if (is_cann_device(src_device) && dst_device.UsesCpuMemory()) {
+        CannToCpuMemCpy(dst_ptr, src_ptr, bytes);
+        return;
+      }
+#endif
+      throw std::runtime_error("Unable to copy data between the source and destination devices");
+    }
+    copy_fn(dst_ptr, src_ptr, bytes);
+  }
+}
+
 }  // namespace python
 }  // namespace onnxruntime

onnxruntime/python/onnxruntime_pybind_mlvalue.h

@@ -138,6 +138,11 @@ pybind11::object GetPyObjFromTensor(const OrtValue& rtensor,
                                     const std::unordered_map<OrtDevice, MemCpyFunc>* mem_cpy_to_host_functions = nullptr,
                                     bool zero_copy_non_owning = false);
 
+// Update the tensor data in an OrtValue in-place from another OrtValue.
+// Both OrtValues must contain tensors of the same data type and size.
+// This function supports various device-to-device transfers.
+void UpdateOrtValueInplace(OrtValue& dst, const OrtValue& src);
+
 // The below two functions are used to convert OrtValue to numpy arrays
 
 /// <summary>

onnxruntime/python/onnxruntime_pybind_ortvalue.cc

@@ -237,6 +237,9 @@ void addOrtValueMethods(pybind11::module& m) {
           throw std::runtime_error("Unsupported device: Cannot update the OrtValue on this device");
         }
       })
+      .def("update_inplace", [](OrtValue* ml_value, const OrtValue& source) {
+        python::UpdateOrtValueInplace(*ml_value, source);
+      })
       // Create an ortvalue value on top of the numpy array, but interpret the data
       // as a different type with the same element size.
       .def_static("ortvalue_from_numpy_with_onnx_type", [](py::array& data, int32_t onnx_element_type) -> std::unique_ptr<OrtValue> {

onnxruntime/test/python/onnxruntime_test_python_cudagraph.py

@@ -76,6 +76,35 @@ def test_ort_value_update_in_place(self):
             ortvalue_gpu.update_inplace(x1)
             np.testing.assert_allclose(ortvalue_gpu.numpy(), x1)
 
+    def test_ort_value_update_in_place_from_ortvalue(self):
+        # Test CPU to CPU copy via OrtValue
+        x0 = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32)
+        x1 = np.array([[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]], dtype=np.float32)
+
+        ortvalue_dst = onnxrt.OrtValue.ortvalue_from_numpy(x0)
+        ortvalue_src = onnxrt.OrtValue.ortvalue_from_numpy(x1)
+        ortvalue_dst.update_inplace(ortvalue_src)
+        np.testing.assert_allclose(ortvalue_dst.numpy(), x1)
+
+        if "CUDAExecutionProvider" in onnxrt.get_available_providers():
+            # Test GPU to GPU copy via OrtValue
+            ortvalue_gpu_dst = onnxrt.OrtValue.ortvalue_from_numpy(x0, "cuda", 0)
+            ortvalue_gpu_src = onnxrt.OrtValue.ortvalue_from_numpy(x1, "cuda", 0)
+            ortvalue_gpu_dst.update_inplace(ortvalue_gpu_src)
+            np.testing.assert_allclose(ortvalue_gpu_dst.numpy(), x1)
+
+            # Test CPU OrtValue to GPU OrtValue copy
+            ortvalue_gpu_dst2 = onnxrt.OrtValue.ortvalue_from_numpy(x0, "cuda", 0)
+            ortvalue_cpu_src = onnxrt.OrtValue.ortvalue_from_numpy(x1)
+            ortvalue_gpu_dst2.update_inplace(ortvalue_cpu_src)
+            np.testing.assert_allclose(ortvalue_gpu_dst2.numpy(), x1)
+
+            # Test GPU OrtValue to CPU OrtValue copy
+            ortvalue_cpu_dst = onnxrt.OrtValue.ortvalue_from_numpy(x0)
+            ortvalue_gpu_src2 = onnxrt.OrtValue.ortvalue_from_numpy(x1, "cuda", 0)
+            ortvalue_cpu_dst.update_inplace(ortvalue_gpu_src2)
+            np.testing.assert_allclose(ortvalue_cpu_dst.numpy(), x1)
+
     def test_select_ep_to_run_cuda_graph(self):
         if "TensorrtExecutionProvider" in onnxrt.get_available_providers():
             providers = [("TensorrtExecutionProvider", {"trt_cuda_graph_enable": True})]