diff --git a/include/onnxruntime/core/framework/ortdevice.h b/include/onnxruntime/core/framework/ortdevice.h index 536d641b4eef9..fea970b84fd84 100644 --- a/include/onnxruntime/core/framework/ortdevice.h +++ b/include/onnxruntime/core/framework/ortdevice.h @@ -150,6 +150,13 @@ struct OrtDevice { return alignment < other.alignment; } + bool EqualIgnoringAlignment(const OrtDevice& other) const { + return device_type == other.device_type && + memory_type == other.memory_type && + vendor_id == other.vendor_id && + device_id == other.device_id; + } + private: // Device type. int32_t device_type : 8; diff --git a/include/onnxruntime/core/session/environment.h b/include/onnxruntime/core/session/environment.h index 306f81df38e48..89467f5238fa9 100644 --- a/include/onnxruntime/core/session/environment.h +++ b/include/onnxruntime/core/session/environment.h @@ -106,6 +106,15 @@ class Environment { return shared_allocators_; } + /** + * Returns an AllocatorPtr for a shared IAllocator based allocator if it matches the memory info. + * The OrtMemoryInfo name and whether it's an arena or device allocator is ignored in the lookup, as is the + * alignment. + * The user calling this function is not expected to know the alignment, and we expect the allocator instance to be + * created with a valid alignment for the device. + */ + AllocatorPtr GetRegisteredSharedAllocator(const OrtMemoryInfo& mem_info) const; + /** * Removes registered allocator that was previously registered for sharing between multiple sessions. */ @@ -171,7 +180,7 @@ class Environment { std::unique_ptr inter_op_thread_pool_; bool create_global_thread_pools_{false}; - std::mutex mutex_; + mutable std::mutex mutex_; // shared allocators from various sources. // CreateAndRegisterAllocator[V2]: IAllocator allocators created by ORT diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index e8d133779f33c..51a8b13cd8261 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -734,6 +734,10 @@ struct CudaEpFactory : OrtEpFactory { } */ + // guard against bad device discovery. max devices we expect to add is num_cuda_devices. if we're attempting + // to add more than that we have duplicates in the `devices` array. + max_ep_devices = std::min(max_ep_devices, static_cast(num_cuda_devices)); + int16_t device_id = 0; for (size_t i = 0; i < num_devices && num_ep_devices < max_ep_devices; ++i) { const OrtHardwareDevice& device = *devices[i]; diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index 2b553aecbca6c..dfb2e33f8cb32 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -72,21 +72,23 @@ ProviderInfo_CUDA& GetProviderInfo_CUDA(); #endif // defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) namespace { -// Ignore whether there is an arena wrapping the allocator by excluding OrtMemoryInfo.alloc_type from the comparison +// Ignore whether there is an arena wrapping the allocator by excluding OrtMemoryInfo.alloc_type from the comparison. static bool AreOrtMemoryInfosEquivalent( const OrtMemoryInfo& left, const OrtMemoryInfo& right, - bool match_name = true) { + bool match_name = true, + bool ignore_alignment = false) { return left.mem_type == right.mem_type && - left.device == right.device && + (ignore_alignment ? left.device.EqualIgnoringAlignment(right.device) : left.device == right.device) && (!match_name || strcmp(left.name, right.name) == 0); } std::vector::const_iterator FindExistingAllocator(const std::vector& allocators, const OrtMemoryInfo& mem_info, - bool match_name = true) { + bool match_name = true, + bool ignore_alignment = false) { auto ite = std::find_if(std::begin(allocators), std::end(allocators), - [&mem_info, match_name](const AllocatorPtr& alloc_ptr) { + [&mem_info, match_name, ignore_alignment](const AllocatorPtr& alloc_ptr) { // We want to do the equality checking of 2 OrtMemoryInfos sans the OrtAllocatorType field. // This is because we want to avoid registering two allocators for the same device that just // differ on OrtAllocatorType. @@ -96,7 +98,8 @@ std::vector::const_iterator FindExistingAllocator(const std::vecto // OrtDeviceAllocator (which is the only accepted value while registering a custom allocator). // If we allowed this, it could potentially cause a lot of confusion as to which shared allocator // to use for that device and we want to avoid having any ugly logic around this. - return AreOrtMemoryInfosEquivalent(alloc_ptr->Info(), mem_info, match_name); + return AreOrtMemoryInfosEquivalent(alloc_ptr->Info(), mem_info, + match_name, ignore_alignment); }); return ite; @@ -428,8 +431,25 @@ Status Environment::CreateAndRegisterAllocatorV2(const std::string& provider_typ } Environment::~Environment() { - // need to make sure all the OrtAllocator instances are released prior to any plugin EPs being freed + // need to make sure all the OrtAllocator instances are released prior to any plugin EPs being freed. + // this is because any entry in shared_allocators_ wrapping an OrtAllocator from a plugin EP owns the OrtAllocator + // instance and will call Release on it. If the plugin EP has been freed the Release will fail. shared_allocators_.clear(); + +#if !defined(ORT_MINIMAL_BUILD) + // unregister any remaining EP libraries so they're cleaned up in a determistic way. + while (!ep_libraries_.empty()) { + auto it = ep_libraries_.begin(); + ORT_IGNORE_RETURN_VALUE(UnregisterExecutionProviderLibrary(it->first)); + } +#endif +} + +AllocatorPtr Environment::GetRegisteredSharedAllocator(const OrtMemoryInfo& mem_info) const { + std::lock_guard lock{mutex_}; + + auto it = FindExistingAllocator(shared_allocators_, mem_info, /*match_name*/ false, /*ignore_alignment*/ true); + return it != shared_allocators_.end() ? *it : nullptr; } Status Environment::GetSharedAllocator(const OrtMemoryInfo& mem_info, OrtAllocator*& allocator) { diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index e8e51db13bcd3..64c4ada07f28f 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -21,7 +21,7 @@ import onnxruntime -def get_ort_device_type(device_type: str, device_index) -> C.OrtDevice: +def get_ort_device_type(device_type: str) -> int: if device_type == "cuda": return C.OrtDevice.cuda() elif device_type == "cann": @@ -32,8 +32,10 @@ def get_ort_device_type(device_type: str, device_index) -> C.OrtDevice: return C.OrtDevice.dml() elif device_type == "webgpu": return C.OrtDevice.webgpu() - elif device_type == "ort": - return C.get_ort_device(device_index).device_type() + elif device_type == "gpu": + return C.OrtDevice.gpu() + elif device_type == "npu": + return C.OrtDevice.npu() else: raise Exception("Unsupported device type: " + device_type) @@ -765,7 +767,7 @@ def bind_input(self, name, device_type, device_id, element_type, shape, buffer_p self._iobinding.bind_input( name, C.OrtDevice( - get_ort_device_type(device_type, device_id), + get_ort_device_type(device_type), C.OrtDevice.default_memory(), device_id, ), @@ -812,7 +814,7 @@ def bind_output( self._iobinding.bind_output( name, C.OrtDevice( - get_ort_device_type(device_type, device_id), + get_ort_device_type(device_type), C.OrtDevice.default_memory(), device_id, ), @@ -823,7 +825,7 @@ def bind_output( self._iobinding.bind_output( name, C.OrtDevice( - get_ort_device_type(device_type, device_id), + get_ort_device_type(device_type), C.OrtDevice.default_memory(), device_id, ), @@ -889,7 +891,7 @@ def _get_c_value(self) -> C.OrtValue: return self._ortvalue @classmethod - def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0) -> OrtValue: + def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0, vendor_id=-1) -> OrtValue: """ Factory method to construct an OrtValue (which holds a Tensor) from a given Numpy object A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu @@ -897,6 +899,7 @@ def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device :param numpy_obj: The Numpy object to construct the OrtValue from :param device_type: e.g. cpu, cuda, cann, cpu by default :param device_id: device id, e.g. 0 + :param vendor_id: The device's PCI vendor id. If provided, the device_type should be "gpu" or "npu". """ # Hold a reference to the numpy object (if device_type is 'cpu') as the OrtValue # is backed directly by the data buffer of the numpy object and so the numpy object @@ -904,11 +907,7 @@ def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device return cls( C.OrtValue.ortvalue_from_numpy( numpy_obj, - C.OrtDevice( - get_ort_device_type(device_type, device_id), - C.OrtDevice.default_memory(), - device_id, - ), + OrtDevice.make(device_type, device_id, vendor_id)._get_c_device(), ), numpy_obj if device_type.lower() == "cpu" else None, ) @@ -929,7 +928,7 @@ def ortvalue_from_numpy_with_onnx_type(cls, data: np.ndarray, /, onnx_element_ty @classmethod def ortvalue_from_shape_and_type( - cls, shape: Sequence[int], element_type, device_type: str = "cpu", device_id: int = 0 + cls, shape: Sequence[int], element_type, device_type: str = "cpu", device_id: int = 0, vendor_id: int = -1 ) -> OrtValue: """ Factory method to construct an OrtValue (which holds a Tensor) from given shape and element_type @@ -938,7 +937,11 @@ def ortvalue_from_shape_and_type( :param element_type: The data type of the elements. It can be either numpy type (like numpy.float32) or an integer for onnx type (like onnx.TensorProto.BFLOAT16). :param device_type: e.g. cpu, cuda, cann, cpu by default :param device_id: device id, e.g. 0 + :param vendor_id: If provided the device type should be "gpu" or "npu". """ + + device = OrtDevice.make(device_type, device_id, vendor_id)._get_c_device() + # Integer for onnx element type (see https://onnx.ai/onnx/api/mapping.html). # This is helpful for some data type (like TensorProto.BFLOAT16) that is not available in numpy. if isinstance(element_type, int): @@ -946,11 +949,7 @@ def ortvalue_from_shape_and_type( C.OrtValue.ortvalue_from_shape_and_onnx_type( shape, element_type, - C.OrtDevice( - get_ort_device_type(device_type, device_id), - C.OrtDevice.default_memory(), - device_id, - ), + device, ) ) @@ -958,11 +957,7 @@ def ortvalue_from_shape_and_type( C.OrtValue.ortvalue_from_shape_and_type( shape, element_type, - C.OrtDevice( - get_ort_device_type(device_type, device_id), - C.OrtDevice.default_memory(), - device_id, - ), + device, ) ) @@ -1085,14 +1080,27 @@ def _get_c_device(self): return self._ort_device @staticmethod - def make(ort_device_name, device_id): - return OrtDevice( - C.OrtDevice( - get_ort_device_type(ort_device_name, device_id), - C.OrtDevice.default_memory(), - device_id, + def make(ort_device_name, device_id, vendor_id=-1): + if vendor_id < 0: + # backwards compatibility with predefined OrtDevice names + return OrtDevice( + C.OrtDevice( + get_ort_device_type(ort_device_name), + C.OrtDevice.default_memory(), + device_id, + ) + ) + else: + # generic. use GPU or NPU for ort_device_name and provide a vendor id. + # vendor id of 0 is valid in some cases (e.g. webgpu is generic and does not have a vendor id) + return OrtDevice( + C.OrtDevice( + get_ort_device_type(ort_device_name), + C.OrtDevice.default_memory(), + vendor_id, + device_id, + ) ) - ) def device_id(self): return self._ort_device.device_id() @@ -1100,6 +1108,9 @@ def device_id(self): def device_type(self): return self._ort_device.device_type() + def device_vendor_id(self): + return self._ort_device.vendor_id() + class SparseTensor: """ diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc index 958c9fc46bcd8..590e1ef3cdbdb 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc @@ -99,6 +99,44 @@ TensorShape GetShape(const py::array& arr) { return shape; } +AllocatorPtr GetSharedAllocator(const OrtDevice& device) { + auto& env = GetOrtEnv()->GetEnvironment(); + + OrtMemoryInfo mem_info("ignored", OrtDeviceAllocator, device); + return env.GetRegisteredSharedAllocator(mem_info); +} + +MemCpyFunc CreateDataTransferMemCpy([[maybe_unused]] const OrtDevice& src_device, + [[maybe_unused]] const OrtDevice& dst_device) { +#if defined(ORT_MINIMAL_BUILD) + // plugin EPs are not supported in a minimal build so there won't be any data transfers registered + return nullptr; +#else + + auto& env = GetOrtEnv()->GetEnvironment(); + const DataTransferManager& data_transfer_manager = env.GetDataTransferManager(); + const IDataTransfer* data_transfer = data_transfer_manager.GetDataTransfer(src_device, dst_device); + if (!data_transfer) { + return nullptr; + } + + const auto copy_func = [src_device, dst_device, data_transfer](void* dst, const void* src, size_t bytes) { + OrtMemoryInfo src_memory_info("ignored", OrtDeviceAllocator, src_device); + OrtMemoryInfo dst_memory_info("ignored", OrtDeviceAllocator, dst_device); + + // real shape doesn't matter as the Tensor instances here are temporary in order to be able to call CopyTensor. + // we set the shape to `bytes` and the data type to uint8_t to copy the correct number of bytes. + TensorShape shape = {narrow(bytes)}; + Tensor src_tensor{DataTypeImpl::GetType(), shape, const_cast(src), src_memory_info}; + Tensor dst_tensor{DataTypeImpl::GetType(), shape, dst, dst_memory_info}; + + ORT_THROW_IF_ERROR(data_transfer->CopyTensor(src_tensor, dst_tensor)); + }; + + return copy_func; +#endif +} + void CpuToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { memcpy(dst, src, num_bytes); } @@ -158,9 +196,10 @@ void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_CUDA().cudaMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetCudaToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, CudaToCpuMemCpy}}; +const std::unordered_map* GetCudaToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, 0}, CudaToCpuMemCpy}, + }; return ↦ } @@ -215,9 +254,10 @@ void MIGraphXToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_MIGraphX().MIGraphXMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetMIGraphXToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, MIGraphXToCpuMemCpy}}; +const std::unordered_map* GetMIGraphXToHostMemCpyFunction(const OrtDevice& device) { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::AMD, 0}, MIGraphXToCpuMemCpy}, + }; return ↦ } @@ -334,9 +374,10 @@ void DmlToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { D3D12_RESOURCE_STATE_UNORDERED_ACCESS); } -const std::unordered_map* GetDmlToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, DmlToCpuMemCpy}}; +const std::unordered_map* GetDmlToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::MICROSOFT, 0}, DmlToCpuMemCpy}, + }; return ↦ } @@ -352,9 +393,10 @@ void CannToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_CANN().cannMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetCannToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::NPU, CannToCpuMemCpy}}; +const std::unordered_map* GetCannToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::NPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::HUAWEI, 0}, CannToCpuMemCpy}, + }; return ↦ } @@ -402,9 +444,10 @@ void RocmToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_ROCM().rocmMemcpy_DeviceToHost(dst, src, num_bytes); } -const std::unordered_map* GetRocmToHostMemCpyFunction() { - static std::unordered_map map{ - {OrtDevice::GPU, RocmToCpuMemCpy}}; +const std::unordered_map* GetRocmToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice{OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::AMD, 0}, RocmToCpuMemCpy}, + }; return ↦ } @@ -581,7 +624,7 @@ using OrtPybindSingleUseAllocatorPtr = std::shared_ptr& p_tensor, - MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) { + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy) { CopyDataToTensor(darray, npy_type, *p_tensor, mem_cpy_to_device); } -void CopyDataToTensor(const py::array& py_array, int npy_type, Tensor& tensor, MemCpyFunc mem_cpy_to_device) { +void CopyDataToTensor(const py::array& py_array, int npy_type, Tensor& tensor, const MemCpyFunc& mem_cpy_to_device) { CopyDataToTensor(reinterpret_cast(py_array.ptr()), npy_type, tensor, mem_cpy_to_device); } @@ -656,7 +699,7 @@ void CopyDataToTensor(const py::array& py_array, int npy_type, Tensor& tensor, M // The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope static std::unique_ptr CreateTensor(const AllocatorPtr& alloc, const std::string& name_input, PyArrayObject* pyObject, bool use_numpy_data_memory = true, - MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) { + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy) { PyArrayObject* darray = PyArray_GETCONTIGUOUS(pyObject); ORT_ENFORCE(darray != nullptr, "The object must be a contiguous array for input '", name_input, "'."); @@ -746,7 +789,8 @@ static void CreateSequenceOfTensors(AllocatorPtr alloc, const std::string& name_ // as the backing data buffer for the ORT Tensor where applicable (for numeric tensors) // The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope static void CreateTensorMLValue(const AllocatorPtr& alloc, const std::string& name_input, PyArrayObject* pyObject, - OrtValue* p_mlvalue, bool use_numpy_data_memory = true, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy) { + OrtValue* p_mlvalue, bool use_numpy_data_memory = true, + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy) { auto p_tensor = CreateTensor(alloc, name_input, pyObject, use_numpy_data_memory, mem_cpy_to_device); auto ml_tensor = DataTypeImpl::GetType(); @@ -994,9 +1038,10 @@ static void CreateGenericIterableMLValue(PyObject* iterator, AllocatorPtr alloc, // Setting `use_numpy_data_memory` to `true` will ensure that the underlying numpy array buffer is directly used // as the backing data buffer for the ORT Tensor where applicable (for numeric tensors) // The numpy object owns the memory and needs to be alive until the corresponding OrtValue is in scope -void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, const std::string& name_input, - const py::object& value, OrtValue* p_mlvalue, bool accept_only_numpy_array, - bool use_numpy_data_memory, MemCpyFunc mem_cpy_to_device) { +void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, + const std::string& name_input, const py::object& value, OrtValue* p_mlvalue, + bool accept_only_numpy_array, bool use_numpy_data_memory, + const MemCpyFunc& mem_cpy_to_device) { onnx::TypeProto type_proto; if (PyObjectCheck_NumpyArray(value.ptr())) { // The most frequent case: input comes as an array. diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.h b/onnxruntime/python/onnxruntime_pybind_mlvalue.h index e9bafea2ed1b5..7b65c0aae45c1 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.h +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.h @@ -42,22 +42,27 @@ MLDataType NumpyTypeToOnnxRuntimeTensorType(int numpy_type); MLDataType OnnxTypeToOnnxRuntimeTensorType(int onnx_element_type); -using MemCpyFunc = void (*)(void*, const void*, size_t); - +using MemCpyFunc = std::function; using DataTransferAlternative = std::variant; +// helpers to get allocator and IDataTransfer from Environment for plugin EP +AllocatorPtr GetSharedAllocator(const OrtDevice& device); +MemCpyFunc CreateDataTransferMemCpy(const OrtDevice& src_device, const OrtDevice& dst_device); + void CpuToCpuMemCpy(void*, const void*, size_t); -void CopyDataToTensor(const pybind11::array& py_array, int npy_type, Tensor& tensor, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy); +void CopyDataToTensor(const pybind11::array& py_array, int npy_type, Tensor& tensor, + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy); pybind11::object AddTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions); + const std::unordered_map* mem_cpy_to_host_functions); -pybind11::object GetPyObjectFromSparseTensor(size_t pos, const OrtValue& ort_value, const DataTransferManager* data_transfer_manager); +pybind11::object GetPyObjectFromSparseTensor(size_t pos, const OrtValue& ort_value, + const DataTransferManager* data_transfer_manager); pybind11::object AddNonTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions); + const std::unordered_map* mem_cpy_to_host_functions); OrtMemoryInfo GetMemoryInfoPerDeviceType(const OrtDevice& ort_device); @@ -69,7 +74,7 @@ void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes); void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetCudaToHostMemCpyFunction(); +const std::unordered_map* GetCudaToHostMemCpyFunction(); bool IsCudaDeviceIdValid(const onnxruntime::logging::Logger& logger, int id); @@ -87,7 +92,7 @@ void CpuToDmlMemCpy(void* dst, const void* src, size_t num_bytes); void DmlToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetDmlToHostMemCpyFunction(); +const std::unordered_map* GetDmlToHostMemCpyFunction(); #endif @@ -97,7 +102,7 @@ void CpuToMIGraphXMemCpy(void* dst, const void* src, size_t num_bytes); void MIGraphXToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetMIGraphXToHostMemCpyFunction(); +const std::unordered_map* GetMIGraphXToHostMemCpyFunction(); AllocatorPtr GetMIGraphXAllocator(OrtDevice::DeviceId id); @@ -109,7 +114,7 @@ void CpuToCannMemCpy(void* dst, const void* src, size_t num_bytes); void CannToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetCannToHostMemCpyFunction(); +const std::unordered_map* GetCannToHostMemCpyFunction(); bool IsCannDeviceIdValid(const onnxruntime::logging::Logger& logger, int id); @@ -127,17 +132,18 @@ void CpuToRocmMemCpy(void* dst, const void* src, size_t num_bytes); void RocmToCpuMemCpy(void* dst, const void* src, size_t num_bytes); -const std::unordered_map* GetRocmToHostMemCpyFunction(); +const std::unordered_map* GetRocmToHostMemCpyFunction(); #endif void CreateGenericMLValue(const onnxruntime::InputDefList* input_def_list, const AllocatorPtr& alloc, const std::string& name_input, const pybind11::object& value, OrtValue* p_mlvalue, - bool accept_only_numpy_array = false, bool use_numpy_data_memory = true, MemCpyFunc mem_cpy_to_device = CpuToCpuMemCpy); + bool accept_only_numpy_array = false, bool use_numpy_data_memory = true, + const MemCpyFunc& mem_cpy_to_device = CpuToCpuMemCpy); pybind11::object GetPyObjFromTensor(const OrtValue& rtensor, const DataTransferManager* data_transfer_manager = nullptr, - const std::unordered_map* mem_cpy_to_host_functions = nullptr); + const std::unordered_map* mem_cpy_to_host_functions = nullptr); // 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 d1d4d6f3cdad5..7234543eb14de 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -23,42 +23,57 @@ std::unique_ptr OrtValueFromShapeAndType(const std::vector& s MLDataType element_type, const OrtDevice& device) { AllocatorPtr allocator; + if (strcmp(GetDeviceName(device), CPU) == 0) { allocator = GetAllocator(); - } else if (strcmp(GetDeviceName(device), CUDA) == 0) { + } else { +#if !defined(ORT_MINIMAL_BUILD) + // prefer a shared allocator from the environment. + // these are provided by plugin EPs or custom allocators explicitly registered by the user. + allocator = GetSharedAllocator(device); +#endif + + if (!allocator) { + if (strcmp(GetDeviceName(device), CUDA) == 0) { #ifdef USE_CUDA - if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { - throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); - } - allocator = GetCudaAllocator(device.Id()); + if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { + throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); + } + + allocator = GetCudaAllocator(device.Id()); #else - throw std::runtime_error( - "Can't allocate memory on the CUDA device using this package of OnnxRuntime. " - "Please use the CUDA package of OnnxRuntime to use this feature."); + throw std::runtime_error( + "Can't allocate memory on the CUDA device using this package of OnnxRuntime. " + "Please use the CUDA package of OnnxRuntime to use this feature."); #endif - } else if (strcmp(GetDeviceName(device), HIP) == 0) { + } else if (strcmp(GetDeviceName(device), HIP) == 0) { #if USE_ROCM - if (!IsRocmDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { - throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); - } - allocator = GetRocmAllocator(device.Id()); + if (!IsRocmDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { + throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); + } + + allocator = GetRocmAllocator(device.Id()); #elif USE_MIGRAPHX - allocator = GetMIGraphXAllocator(device.Id()); + allocator = GetMIGraphXAllocator(device.Id()); #else - throw std::runtime_error( - "Can't allocate memory on the AMD device using this package of OnnxRuntime. " - "Please use the ROCm package of OnnxRuntime to use this feature."); + throw std::runtime_error( + "Can't allocate memory on the AMD device using this package of OnnxRuntime. " + "Please use the ROCm package of OnnxRuntime to use this feature."); #endif - } else if (strcmp(GetDeviceName(device), DML) == 0) { + } else if (strcmp(GetDeviceName(device), DML) == 0) { #if USE_DML - allocator = GetDmlAllocator(device.Id()); + allocator = GetDmlAllocator(device.Id()); #else - throw std::runtime_error( - "Can't allocate memory on the DirectML device using this package of OnnxRuntime. " - "Please use the DirectML package of OnnxRuntime to use this feature."); + throw std::runtime_error( + "Can't allocate memory on the DirectML device using this package of OnnxRuntime. " + "Please use the DirectML package of OnnxRuntime to use this feature."); #endif - } else { - throw std::runtime_error("Unsupported device: Cannot place the OrtValue on this device"); + } + } + + if (!allocator) { + throw std::runtime_error("Unsupported device: Cannot place the OrtValue on this device"); + } } auto ml_value = std::make_unique(); @@ -90,7 +105,8 @@ void addOrtValueMethods(pybind11::module& m) { if (device.Vendor() == OrtDevice::VendorIds::MICROSOFT) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in DML + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in DML CreateGenericMLValue( nullptr, GetDmlAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToDmlMemCpy); } else @@ -103,8 +119,10 @@ void addOrtValueMethods(pybind11::module& m) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in CUDA - CreateGenericMLValue(nullptr, GetCudaAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToCudaMemCpy); + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in CUDA + CreateGenericMLValue(nullptr, GetCudaAllocator(device.Id()), "", array_on_cpu, ml_value.get(), + true, false, CpuToCudaMemCpy); } else #endif #ifdef USE_ROCM @@ -115,22 +133,34 @@ void addOrtValueMethods(pybind11::module& m) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in CUDA - CreateGenericMLValue(nullptr, GetRocmAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToRocmMemCpy); + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in ROCM + CreateGenericMLValue(nullptr, GetRocmAllocator(device.Id()), "", array_on_cpu, ml_value.get(), + true, false, CpuToRocmMemCpy); } else #endif #if USE_MIGRAPHX if (device.Vendor() == OrtDevice::VendorIds::AMD) { // InputDeflist is null because OrtValue creation is not tied to a specific model // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) - // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors in MIGraphX - CreateGenericMLValue(nullptr, GetMIGraphXAllocator(device.Id()), "", array_on_cpu, ml_value.get(), true, false, CpuToMIGraphXMemCpy); + // TODO: Add check to ensure that string arrays are not passed - we currently don't support string tensors + // in MIGraphX + CreateGenericMLValue(nullptr, GetMIGraphXAllocator(device.Id()), "", array_on_cpu, ml_value.get(), + true, false, CpuToMIGraphXMemCpy); } else #endif { - throw std::runtime_error( - "Can't allocate memory on the CUDA device using this package of OnnxRuntime. " - "Please use the CUDA package of OnnxRuntime to use this feature."); + // see if we can do the copy with an allocator and IDataTransfer registered by a plugin EP + auto allocator = GetSharedAllocator(device); + auto cpu_to_device_copy_fn = allocator ? CreateDataTransferMemCpy(OrtDevice{}, device) : nullptr; + if (cpu_to_device_copy_fn) { + CreateGenericMLValue(nullptr, allocator, "", array_on_cpu, ml_value.get(), true, false, + cpu_to_device_copy_fn); + } else { + throw std::runtime_error( + "Can't allocate memory on the device using this package of OnnxRuntime. " + "Please use the appropriate package of OnnxRuntime for your hardware to use this feature."); + } } } else if (device.Type() == OrtDevice::NPU && device.Vendor() == OrtDevice::VendorIds::HUAWEI) { #ifdef USE_CANN @@ -214,8 +244,16 @@ void addOrtValueMethods(pybind11::module& m) { } else #endif { - throw std::runtime_error( - "Unsupported GPU device: Cannot find the supported GPU device."); + // see if we can do the copy with an allocator and IDataTransfer registered by a plugin EP + auto allocator = GetSharedAllocator(device); + auto cpu_to_device_copy_fn = allocator ? CreateDataTransferMemCpy(OrtDevice{}, device) : nullptr; + if (cpu_to_device_copy_fn) { + onnxruntime::python::CopyDataToTensor(py_values, values_type, *(ml_value->GetMutable()), + cpu_to_device_copy_fn); + } else { + throw std::runtime_error( + "Unsupported GPU device: Cannot find the supported GPU device."); + } } } else if (device.Type() == OrtDevice::DML) { #if USE_DML diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index acf0681cf8752..03ad0185d1394 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -205,7 +205,7 @@ void AppendLoraParametersAsInputs(const RunOptions& run_options, template static py::object AddNonTensor(const OrtValue& val, const DataTransferManager* /*data_transfer_manager*/, - const std::unordered_map* /*mem_cpy_to_host_functions*/) { + const std::unordered_map* /*mem_cpy_to_host_functions*/) { return py::cast(val.Get()); } @@ -265,39 +265,65 @@ pybind11::array PrimitiveTensorToNumpyFromDevice(const OrtValue& ort_value, cons // pretty much does what a DataTransferManager does - copy data from device(s) to the host py::object GetPyObjFromTensor(const OrtValue& ort_value, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { ORT_ENFORCE(ort_value.IsTensor(), "This function only supports tensors"); const auto& tensor = ort_value.Get(); + const auto& device = tensor.Location().device; + if (tensor.IsDataTypeString()) { - ORT_ENFORCE(tensor.Location().device.Type() == OrtDevice::CPU, "Strings can only be on CPU"); + ORT_ENFORCE(device.Type() == OrtDevice::CPU, "Strings can only be on CPU"); // Create a numpy array of strings (python objects) by copy/converting them py::array result = StringTensorToNumpyArray(tensor); return py::cast(result); } - const auto device_type = tensor.Location().device.Type(); + const auto device_type = device.Type(); // Create an numpy array on top of the OrtValue memory, no copy if (device_type == OrtDevice::CPU) { py::array result = PrimitiveTensorToNumpyOverOrtValue(ort_value); return py::cast(result); } - if (!data_transfer_manager && !mem_cpy_to_host_functions) { - throw std::runtime_error( - "GetPyObjFromTensor: Either data transfer manager or a " - "function to copy data to the host is needed to convert non-CPU tensor to numpy array"); - } - py::array result; if (data_transfer_manager != nullptr) { result = PrimitiveTensorToNumpyFromDevice(ort_value, data_transfer_manager); } else { - auto mem_cpy_to_host = mem_cpy_to_host_functions->find(device_type); - ORT_ENFORCE(mem_cpy_to_host != mem_cpy_to_host_functions->end(), - "Unable to locate a function that can copy data to the host from the device"); - result = PrimitiveTensorToNumpyFromDevice(ort_value, mem_cpy_to_host->second); + bool copied = false; + if (mem_cpy_to_host_functions) { + auto it = std::find_if(mem_cpy_to_host_functions->begin(), mem_cpy_to_host_functions->end(), + [&device](const auto& entry) { + const auto& copy_device = entry.first; + // We're ignoring OrtDevice.Id() currently for historical reasons. + // The key to mem_cpy_to_host_functions was previously the device type (CPU/GPU/NPU). + // This changed to be OrtDevice to get the vendor id. + // Assumably it would be better to also match on device id, but that was not possible + // previously and to preserve existing behavior we keep the old logic and expect the + // copy function to handle the device id correctly. + return device.Type() == copy_device.Type() && + device.MemType() == copy_device.MemType() && + device.Vendor() == copy_device.Vendor(); + }); + + if (it != mem_cpy_to_host_functions->end()) { + result = PrimitiveTensorToNumpyFromDevice(ort_value, it->second); + copied = true; + } + } + + if (!copied) { + // see if we have a shared data transfer function from a plugin EP + auto device_to_cpu_copy_func = CreateDataTransferMemCpy(device, OrtDevice{}); + if (device_to_cpu_copy_func) { + result = PrimitiveTensorToNumpyFromDevice(ort_value, device_to_cpu_copy_func); + } else { + throw std::runtime_error( + "GetPyObjFromTensor: Either data transfer manager or a " + "function to copy data to the host is needed to convert non-CPU tensor to numpy array"); + } + } } + return py::cast(result); } @@ -373,7 +399,7 @@ py::object GetPyObjectFromSparseTensor(size_t pos, const OrtValue& ort_value, co template <> py::object AddNonTensor(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { const auto& seq_tensors = val.Get(); py::list py_list; for (const auto& ort_value : seq_tensors) { @@ -389,7 +415,7 @@ py::object AddNonTensor(const OrtValue& val, py::object AddNonTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { // Should be in sync with core/framework/datatypes.h auto val_type = val.Type(); if (val_type->IsTensorSequenceType()) { @@ -429,7 +455,7 @@ py::object AddNonTensorAsPyObj(const OrtValue& val, } py::object AddTensorAsPyObj(const OrtValue& val, const DataTransferManager* data_transfer_manager, - const std::unordered_map* mem_cpy_to_host_functions) { + const std::unordered_map* mem_cpy_to_host_functions) { return GetPyObjFromTensor(val, data_transfer_manager, mem_cpy_to_host_functions); } @@ -1885,6 +1911,10 @@ void addObjectMethods(py::module& m, ExecutionProviderRegistrationFn ep_registra vendor = OrtDevice::VendorIds::NVIDIA; #elif USE_ROCM || USE_MIGRAPHX vendor = OrtDevice::VendorIds::AMD; +#endif + } else if (type == OrtDevice::NPU) { +#if USE_CANN + vendor = OrtDevice::VendorIds::HUAWEI; #endif } @@ -1894,12 +1924,15 @@ void addObjectMethods(py::module& m, ExecutionProviderRegistrationFn ep_registra .def("device_id", &OrtDevice::Id, R"pbdoc(Device Id.)pbdoc") .def("device_type", &OrtDevice::Type, R"pbdoc(Device Type.)pbdoc") .def("vendor_id", &OrtDevice::Vendor, R"pbdoc(Vendor Id.)pbdoc") + // generic device types that are typically used with a vendor id. .def_static("cpu", []() { return OrtDevice::CPU; }) + .def_static("gpu", []() { return OrtDevice::GPU; }) + .def_static("npu", []() { return OrtDevice::NPU; }) + // EP specific device types for backward compatibility. .def_static("cuda", []() { return OrtDevice::GPU; }) .def_static("cann", []() { return OrtDevice::NPU; }) - .def_static("fpga", []() { return OrtDevice::FPGA; }) - .def_static("npu", []() { return OrtDevice::NPU; }) .def_static("dml", []() { return OrtDevice::DML; }) + .def_static("fpga", []() { return OrtDevice::FPGA; }) .def_static("webgpu", []() { return OrtDevice::GPU; }) .def_static("default_memory", []() { return OrtDevice::MemType::DEFAULT; }); diff --git a/onnxruntime/test/python/onnxruntime_test_python_autoep.py b/onnxruntime/test/python/onnxruntime_test_python_autoep.py index 0c52740398b7a..cb31627a87c48 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_autoep.py +++ b/onnxruntime/test/python/onnxruntime_test_python_autoep.py @@ -183,7 +183,7 @@ def test_example_plugin_ep_devices(self): Test registration of an example EP plugin and retrieval of its OrtEpDevice. """ if sys.platform != "win32": - self.skipTest("Skipping test because it device discovery is only supported on Windows") + self.skipTest("Skipping test because device discovery is only supported on Windows") ep_lib_path = "example_plugin_ep.dll" try: @@ -244,6 +244,44 @@ def test_example_plugin_ep_devices(self): del sess # Delete session before unregistering library self.unregister_execution_provider_library(ep_name) + def test_example_plugin_ep_data_transfer(self): + """ + Test usage of shared data transfer and allocator from plugin EP. + """ + if sys.platform != "win32": + self.skipTest("Skipping test because device discovery is only supported on Windows") + + if "DmlExecutionProvider" in onnxrt.get_available_providers(): + self.skipTest("Skipping because DML EP data transfer is broken if we haven't created an inference session") + + ep_lib_path = "example_plugin_ep.dll" + try: + ep_lib_path = get_name("example_plugin_ep.dll") + except FileNotFoundError: + self.skipTest(f"Skipping test because EP library '{ep_lib_path}' cannot be found") + + ep_name = "example_ep" + self.register_execution_provider_library(ep_name, os.path.realpath(ep_lib_path)) + + data = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) + data2 = data + 1 + + # the example EP pretends to use GPU memory so we can test data transfer. + # by matching its OrtDevice info we will hit its allocator and data transfer implementations. + # copy data from CPU to the fake GPU memory + gpu_value = onnxrt.OrtValue.ortvalue_from_numpy(data, "gpu", 0, 0xBE57) + # copy back to CPU + cpu_data = gpu_value.numpy() + np.testing.assert_equal(data, cpu_data) + + gpu_value.update_inplace(data2) # update the fake GPU data + cpu_data_2 = gpu_value.numpy() # copy back to CPU + np.testing.assert_equal(data2, cpu_data_2) + + gpu_value = None # Delete OrtValue before unregistering library as the allocator will be destroyed. + + self.unregister_execution_provider_library(ep_name) + if __name__ == "__main__": unittest.main(verbosity=1)