diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index def2240358c10..e35e3c5753d36 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/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: diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc index 89651c2d955de..fa609fe6ea83d 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc @@ -1071,5 +1071,116 @@ 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(); + const auto& src_tensor = src.Get(); + + 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()->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. +#ifdef USE_CUDA + const auto is_cuda_device = [](const OrtDevice& device) { + return device.Type() == OrtDevice::GPU && device.Vendor() == OrtDevice::VendorIds::NVIDIA; + }; + + 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())); + 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 + const auto is_migraphx_device = [](const OrtDevice& device) { + return device.Type() == OrtDevice::GPU && device.Vendor() == OrtDevice::VendorIds::AMD; + }; + + 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 + const auto is_dml_device = [](const OrtDevice& device) { + return (device.Type() == OrtDevice::GPU && device.Vendor() == OrtDevice::VendorIds::MICROSOFT) || + device.Type() == OrtDevice::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 + const auto is_cann_device = [](const OrtDevice& device) { + return device.Type() == OrtDevice::NPU && device.Vendor() == OrtDevice::VendorIds::HUAWEI; + }; + + 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 diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.h b/onnxruntime/python/onnxruntime_pybind_mlvalue.h index 144b3edcad404..097c5b4d20d65 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.h +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.h @@ -138,6 +138,11 @@ pybind11::object GetPyObjFromTensor(const OrtValue& rtensor, const std::unordered_map* 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 /// diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index eb966ac5fc314..168d57fc0827b 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/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 { diff --git a/onnxruntime/test/python/onnxruntime_test_python_cudagraph.py b/onnxruntime/test/python/onnxruntime_test_python_cudagraph.py index d6c1dd9cff3f3..987efd5af5e8e 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_cudagraph.py +++ b/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})]