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

onnxruntime/python/onnxruntime_inference_collection.py

@@ -1302,15 +1302,21 @@ 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)
+        else:
+            self._ortvalue.update_inplace(data)
 
 
 def copy_tensors(src: Sequence[OrtValue], dst: Sequence[OrtValue], stream=None) -> None:

onnxruntime/python/onnxruntime_pybind_ortvalue.cc

@@ -237,6 +237,97 @@ 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) {
+        if (!ml_value->IsTensor()) {
+          throw std::runtime_error("Inplace update of OrtValues is only supported for Tensors");
+        }
+        if (!source->IsTensor()) {
+          throw std::runtime_error("The source OrtValue must contain a Tensor");
+        }
+
+        const auto& dst_tensor = ml_value->Get<Tensor>();
+        const auto& src_tensor = source->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 = ml_value->GetMutable<Tensor>()->MutableDataRaw();
+        const void* src = src_tensor.DataRaw();
+
+        if (src_device.UsesCpuMemory() && dst_device.UsesCpuMemory()) {
+          memcpy(dst, src, bytes);
+        } else {
+          auto copy_fn = CreateDataTransferMemCpy(src_device, dst_device);
+          if (!copy_fn) {
+            // Fall back to built-in EP copy functions
+#ifdef USE_CUDA
+            if (src_device.Type() == OrtDevice::GPU && dst_device.Type() == OrtDevice::GPU) {
+              auto data_transfer = GetGPUDataTransfer();
+              ORT_THROW_IF_ERROR(data_transfer->CopyTensor(src_tensor, *(ml_value->GetMutable<Tensor>())));
+              return;
+            }
+            if (src_device.UsesCpuMemory() && dst_device.Type() == OrtDevice::GPU) {
+              CpuToCudaMemCpy(dst, src, bytes);
+              return;
+            }
+            if (src_device.Type() == OrtDevice::GPU && dst_device.UsesCpuMemory()) {
+              CudaToCpuMemCpy(dst, src, bytes);
+              return;
+            }
+#endif
+#if USE_MIGRAPHX
+            if (src_device.UsesCpuMemory() && dst_device.Type() == OrtDevice::GPU) {
+              CpuToMIGraphXMemCpy(dst, src, bytes);
+              return;
+            }
+            if (src_device.Type() == OrtDevice::GPU && dst_device.UsesCpuMemory()) {
+              MIGraphXToCpuMemCpy(dst, src, bytes);
+              return;
+            }
+#endif
+#if USE_DML
+            if (src_device.UsesCpuMemory() && (dst_device.Type() == OrtDevice::GPU || dst_device.Type() == OrtDevice::DML)) {
+              CpuToDmlMemCpy(dst, src, bytes);
+              return;
+            }
+            if ((src_device.Type() == OrtDevice::GPU || src_device.Type() == OrtDevice::DML) && dst_device.UsesCpuMemory()) {
+              DmlToCpuMemCpy(dst, src, bytes);
+              return;
+            }
+#endif
+#ifdef USE_CANN
+            if (src_device.UsesCpuMemory() && dst_device.Type() == OrtDevice::NPU) {
+              CpuToCannMemCpy(dst, src, bytes);
+              return;
+            }
+            if (src_device.Type() == OrtDevice::NPU && dst_device.UsesCpuMemory()) {
+              CannToCpuMemCpy(dst, src, bytes);
+              return;
+            }
+#endif
+            throw std::runtime_error("Unable to copy data between the source and destination devices");
+          }
+          copy_fn(dst, src, bytes);
+        }
+      })
       // 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})]