From 38c2f759d497a3125017915e5564e400e96bf685 Mon Sep 17 00:00:00 2001 From: Selena Yang <179177246+selenayang888@users.noreply.github.com> Date: Fri, 16 Oct 2026 21:55:38 -0700 Subject: [PATCH 1/4] Add EP device memory capacity query API for total and available memory --- .../core/session/onnxruntime_c_api.h | 27 ++++ .../core/session/onnxruntime_cxx_api.h | 1 + .../core/session/onnxruntime_cxx_inline.h | 8 + .../core/session/onnxruntime_ep_c_api.h | 25 +++ .../providers/cuda/cuda_provider_factory.cc | 34 +++- .../providers/cuda/plugin/cuda_ep_factory.cc | 31 ++++ .../providers/cuda/plugin/cuda_ep_factory.h | 6 + .../cuda/shared_inc/cuda_memory_capacity.h | 27 ++++ .../cuda/shared_inc/device_memory_capacity.h | 55 +++++++ onnxruntime/core/session/onnxruntime_c_api.cc | 29 ++++ onnxruntime/core/session/ort_apis.h | 2 + .../session/plugin_ep/ep_factory_internal.cc | 1 + .../session/plugin_ep/ep_factory_internal.h | 6 + .../plugin_ep/ep_factory_internal_impl.h | 8 + .../plugin_ep/ep_factory_provider_bridge.h | 12 ++ .../plugin_ep/forward_to_factory_impl.h | 9 ++ .../library/example_plugin_ep/ep_factory.cc | 10 ++ .../library/example_plugin_ep/ep_factory.h | 6 + onnxruntime/test/autoep/test_registration.cc | 146 ++++++++++++++++++ 19 files changed, 442 insertions(+), 1 deletion(-) create mode 100644 onnxruntime/core/providers/cuda/shared_inc/cuda_memory_capacity.h create mode 100644 onnxruntime/core/providers/cuda/shared_inc/device_memory_capacity.h diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 572dbe2aa0a22..63fa51838da58 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -7706,6 +7706,33 @@ struct OrtApi { * \since Version 1.31. */ const OrtModelPackageApi*(ORT_API_CALL* GetModelPackageApi)(void); + + /** \brief Get the current device memory capacity for an execution provider device. + * + * The execution provider queries the device when this function is called. `total_memory_bytes` + * reports device-global memory capacity exposed by the provider for this device, not a per-process + * or per-context allocation budget. For partitioned or virtual devices, this is the capacity of + * the exposed device, not necessarily the entire physical board. + * `available_memory_bytes` is a current free-memory estimate, not a reservation or a guarantee that + * an allocation of that size will succeed. It may change immediately after this call, and allocator + * caches or process limits may further constrain allocations. + * + * The query may initialize provider resources. For example, CUDA may initialize the device's primary + * context to query free memory, while preserving the calling thread's previous current CUDA context. + * + * If the execution provider does not support this query, ORT_NOT_IMPLEMENTED is returned. + * On failure, the output values are not modified. + * + * \param[in] ep_device The OrtEpDevice instance to query. + * \param[out] total_memory_bytes The device-global memory capacity in bytes. + * \param[out] available_memory_bytes The current free device memory estimate in bytes. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.32. + */ + ORT_API2_STATUS(EpDevice_GetMemoryCapacity, _In_ const OrtEpDevice* ep_device, + _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes); }; /* diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index e7bed2238ac22..7917b8e1a2f0e 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -1186,6 +1186,7 @@ struct EpDeviceImpl : Ort::detail::Base { ConstKeyValuePairs EpOptions() const; ConstHardwareDevice Device() const; ConstMemoryInfo GetMemoryInfo(OrtDeviceMemoryType memory_type) const; ///< Wraps EpDevice_MemoryInfo + std::pair GetMemoryCapacity() const; ///< Total and available device memory SyncStream CreateSyncStream(ConstKeyValuePairs stream_options = {}) const; /// Wraps EpDevice_CreateSyncStream }; } // namespace detail diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index a7a9dd4dc5f86..e838355539297 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -756,6 +756,14 @@ inline ConstMemoryInfo EpDeviceImpl::GetMemoryInfo(OrtDeviceMemoryType memory return ConstMemoryInfo{mem_info}; } +template +inline std::pair EpDeviceImpl::GetMemoryCapacity() const { + uint64_t total_memory_bytes{}; + uint64_t available_memory_bytes{}; + ThrowOnError(GetApi().EpDevice_GetMemoryCapacity(this->p_, &total_memory_bytes, &available_memory_bytes)); + return {total_memory_bytes, available_memory_bytes}; +} + template inline SyncStream EpDeviceImpl::CreateSyncStream(ConstKeyValuePairs stream_options) const { OrtSyncStream* stream = nullptr; diff --git a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h index 77a1b7c0f1683..46f3493126119 100644 --- a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h @@ -3308,6 +3308,31 @@ struct OrtEpFactory { _In_ size_t num_candidates, _In_opt_ const OrtSessionOptions* session_options, _Out_ size_t* selected_index); + + /** \brief Get the current device memory capacity for an OrtEpDevice. + * + * This function is optional. Implementations must query the device when called because available + * memory is dynamic. If the query is unsupported, leave this callback null or return ORT_NOT_IMPLEMENTED. + * Total memory must be the device-global capacity exposed for this device, not a per-process or + * per-context allocation budget. For partitioned or virtual devices, report the exposed device's + * capacity, not the entire physical board's capacity. Available memory is a current free-memory + * estimate, not a reservation or an allocation guarantee; allocator caches and process limits may + * further constrain allocations. Implementations must preserve the calling thread's current device + * context if they temporarily switch contexts to perform the query. + * On failure, output values must not be modified. + * + * \param[in] this_ptr The OrtEpFactory instance. + * \param[in] ep_device The OrtEpDevice instance to query. + * \param[out] total_memory_bytes The device-global memory capacity in bytes. + * \param[out] available_memory_bytes The current free device memory estimate in bytes. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.32. + */ + ORT_API2_STATUS(GetDeviceMemoryCapacity, _In_ OrtEpFactory* this_ptr, + _In_ const OrtEpDevice* ep_device, + _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes); }; #ifdef __cplusplus diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 00519584950b8..2a8bbaf3586db 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -19,6 +19,7 @@ #include "core/providers/cuda/cuda_allocator.h" #include "core/providers/cuda/cuda_stream_handle.h" #include "core/providers/cuda/gpu_data_transfer.h" +#include "core/providers/cuda/shared_inc/cuda_memory_capacity.h" #include "core/providers/cuda/math/unary_elementwise_ops_impl.h" #ifdef ENABLE_NVTX_PROFILE @@ -688,10 +689,12 @@ struct CudaSyncStreamImpl : OrtSyncStreamImpl { struct CudaEpFactory : OrtEpFactory { using MemoryInfoUniquePtr = std::unique_ptr>; - CudaEpFactory(const OrtApi& ort_api_in, const OrtLogger& default_logger_in) : ort_api{ort_api_in}, + CudaEpFactory(const OrtApi& ort_api_in, const OrtLogger& default_logger_in) : OrtEpFactory{}, + ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()}, default_logger{default_logger_in}, data_transfer_impl{ort_api_in} { + ort_version_supported = ORT_API_VERSION; GetName = GetNameImpl; GetVendor = GetVendorImpl; GetVendorId = GetVendorIdImpl; @@ -706,6 +709,7 @@ struct CudaEpFactory : OrtEpFactory { IsStreamAware = IsStreamAwareImpl; CreateSyncStreamForDevice = CreateSyncStreamForDeviceImpl; + GetDeviceMemoryCapacity = GetDeviceMemoryCapacityImpl; } static const char* GetNameImpl(const OrtEpFactory* this_ptr) noexcept { @@ -903,6 +907,34 @@ struct CudaEpFactory : OrtEpFactory { return nullptr; } + static OrtStatus* ORT_API_CALL GetDeviceMemoryCapacityImpl( + OrtEpFactory* this_ptr, const OrtEpDevice* ep_device, + uint64_t* total_memory_bytes, uint64_t* available_memory_bytes) noexcept { + auto& factory = *static_cast(this_ptr); + const OrtMemoryInfo* memory_info = factory.ort_api.EpDevice_MemoryInfo( + ep_device, OrtDeviceMemoryType_DEFAULT); + if (memory_info == nullptr) { + return factory.ort_api.CreateStatus(ORT_FAIL, "CUDA device memory info is unavailable"); + } + + int device_id = 0; + OrtStatus* status = factory.ort_api.MemoryInfoGetId(memory_info, &device_id); + if (status != nullptr) { + return status; + } + + cuda::MemoryCapacityDriverApi driver; + uint64_t total{}; + uint64_t available{}; + if (const char* error = cuda::QueryDeviceMemoryCapacity(driver, device_id, total, available)) { + return factory.ort_api.CreateStatus(ORT_RUNTIME_EXCEPTION, error); + } + + *total_memory_bytes = total; + *available_memory_bytes = available; + return nullptr; + } + OrtStatus* CreateMemoryInfoForDevices(int num_devices) { gpu_memory_infos.reserve(num_devices); host_accessible_memory_infos.reserve(num_devices); diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc index 3e4a7bb3c0332..08a87f1e40f44 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc @@ -5,6 +5,7 @@ #include "cuda_device_mapping.h" #include "cuda_ep.h" #include "cuda_plugin_kernels.h" +#include "core/providers/cuda/shared_inc/cuda_memory_capacity.h" #include "core/common/string_utils.h" #include "core/session/onnxruntime_c_api.h" #include @@ -51,6 +52,7 @@ CudaEpFactory::CudaEpFactory(const OrtApi& ort_api, const OrtEpApi& ep_api, CreateDataTransfer = CreateDataTransferImpl; IsStreamAware = IsStreamAwareImpl; CreateSyncStreamForDevice = CreateSyncStreamForDeviceImpl; + GetDeviceMemoryCapacity = GetDeviceMemoryCapacityImpl; } CudaEpFactory::~CudaEpFactory() { @@ -115,6 +117,35 @@ const char* ORT_API_CALL CudaEpFactory::GetVersionImpl(const OrtEpFactory* this_ return static_cast(this_ptr)->ep_version_.c_str(); } +/*static*/ +OrtStatus* ORT_API_CALL CudaEpFactory::GetDeviceMemoryCapacityImpl( + OrtEpFactory* this_ptr, const OrtEpDevice* ep_device, + uint64_t* total_memory_bytes, uint64_t* available_memory_bytes) noexcept { + auto* factory = static_cast(this_ptr); + const OrtMemoryInfo* memory_info = factory->ort_api_.EpDevice_MemoryInfo( + ep_device, OrtDeviceMemoryType_DEFAULT); + if (memory_info == nullptr) { + return factory->ort_api_.CreateStatus(ORT_FAIL, "CUDA device memory info is unavailable"); + } + + int device_id = 0; + OrtStatus* status = factory->ort_api_.MemoryInfoGetId(memory_info, &device_id); + if (status != nullptr) { + return status; + } + + cuda::MemoryCapacityDriverApi driver; + uint64_t total{}; + uint64_t available{}; + if (const char* error = cuda::QueryDeviceMemoryCapacity(driver, device_id, total, available)) { + return factory->ort_api_.CreateStatus(ORT_RUNTIME_EXCEPTION, error); + } + + *total_memory_bytes = total; + *available_memory_bytes = available; + return nullptr; +} + namespace { std::string ToUpper(std::string value) { diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h index dec2d7b1e5992..cc1bafe038aea 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h @@ -98,6 +98,12 @@ class CudaEpFactory : public OrtEpFactory { const OrtKeyValuePairs* stream_options, OrtSyncStreamImpl** stream) noexcept; + static OrtStatus* ORT_API_CALL GetDeviceMemoryCapacityImpl( + OrtEpFactory* this_ptr, + const OrtEpDevice* ep_device, + uint64_t* total_memory_bytes, + uint64_t* available_memory_bytes) noexcept; + const OrtApi& ort_api_; const OrtEpApi& ep_api_; const OrtLogger& default_logger_; diff --git a/onnxruntime/core/providers/cuda/shared_inc/cuda_memory_capacity.h b/onnxruntime/core/providers/cuda/shared_inc/cuda_memory_capacity.h new file mode 100644 index 0000000000000..6e0e84d2874bb --- /dev/null +++ b/onnxruntime/core/providers/cuda/shared_inc/cuda_memory_capacity.h @@ -0,0 +1,27 @@ +#pragma once + +#include + +#include "device_memory_capacity.h" + +namespace onnxruntime { +namespace cuda { + +struct MemoryCapacityDriverApi { + using Device = CUdevice; + using Context = CUcontext; + + bool Init() noexcept { return cuInit(0) == CUDA_SUCCESS; } + bool DeviceGet(Device* device, int ordinal) noexcept { return cuDeviceGet(device, ordinal) == CUDA_SUCCESS; } + bool DeviceTotalMem(size_t* total, Device device) noexcept { return cuDeviceTotalMem(total, device) == CUDA_SUCCESS; } + bool PrimaryCtxRetain(Context* context, Device device) noexcept { + return cuDevicePrimaryCtxRetain(context, device) == CUDA_SUCCESS; + } + bool CtxPushCurrent(Context context) noexcept { return cuCtxPushCurrent(context) == CUDA_SUCCESS; } + bool MemGetInfo(size_t* available, size_t* total) noexcept { return cuMemGetInfo(available, total) == CUDA_SUCCESS; } + bool CtxPopCurrent(Context* context) noexcept { return cuCtxPopCurrent(context) == CUDA_SUCCESS; } + bool PrimaryCtxRelease(Device device) noexcept { return cuDevicePrimaryCtxRelease(device) == CUDA_SUCCESS; } +}; + +} // namespace cuda +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/cuda/shared_inc/device_memory_capacity.h b/onnxruntime/core/providers/cuda/shared_inc/device_memory_capacity.h new file mode 100644 index 0000000000000..e2fb04f6f016d --- /dev/null +++ b/onnxruntime/core/providers/cuda/shared_inc/device_memory_capacity.h @@ -0,0 +1,55 @@ +#pragma once + +#include +#include + +namespace onnxruntime { +namespace cuda { + +template +const char* QueryDeviceMemoryCapacity(DriverApi& driver, int device_id, + uint64_t& total_memory_bytes, + uint64_t& available_memory_bytes) noexcept { + typename DriverApi::Device device{}; + typename DriverApi::Context context{}; + if (!driver.Init()) { + return "cuInit failed"; + } + if (!driver.DeviceGet(&device, device_id)) { + return "cuDeviceGet failed"; + } + size_t total{}; + if (!driver.DeviceTotalMem(&total, device)) { + return "cuDeviceTotalMem failed"; + } + if (!driver.PrimaryCtxRetain(&context, device)) { + return "cuDevicePrimaryCtxRetain failed"; + } + if (!driver.CtxPushCurrent(context)) { + driver.PrimaryCtxRelease(device); + return "cuCtxPushCurrent failed"; + } + + size_t available{}; + size_t context_total{}; + const bool query_ok = driver.MemGetInfo(&available, &context_total); + typename DriverApi::Context popped_context{}; + const bool restore_ok = driver.CtxPopCurrent(&popped_context); + const bool release_ok = driver.PrimaryCtxRelease(device); + if (!restore_ok) { + return "cuCtxPopCurrent failed to restore the calling thread's CUDA context"; + } + if (!release_ok) { + return "cuDevicePrimaryCtxRelease failed"; + } + if (!query_ok) { + return "cuMemGetInfo failed"; + } + + total_memory_bytes = static_cast(total); + available_memory_bytes = static_cast(available); + return nullptr; +} + +} // namespace cuda +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index 84ba4abbc9f79..3541c4aae6ee9 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -4310,6 +4310,34 @@ ORT_API(const OrtMemoryInfo*, OrtApis::EpDevice_MemoryInfo, _In_ const OrtEpDevi } } +ORT_API_STATUS_IMPL(OrtApis::EpDevice_GetMemoryCapacity, _In_ const OrtEpDevice* ep_device, + _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes) { + API_IMPL_BEGIN + ORT_API_RETURN_IF(ep_device == nullptr, ORT_INVALID_ARGUMENT, "ep_device must not be null"); + ORT_API_RETURN_IF(total_memory_bytes == nullptr, ORT_INVALID_ARGUMENT, + "total_memory_bytes must not be null"); + ORT_API_RETURN_IF(available_memory_bytes == nullptr, ORT_INVALID_ARGUMENT, + "available_memory_bytes must not be null"); + + OrtEpFactory* factory = ep_device->GetMutableFactory(); + ORT_API_RETURN_IF(factory == nullptr || factory->ort_version_supported < 32 || + factory->GetDeviceMemoryCapacity == nullptr, + ORT_NOT_IMPLEMENTED, + "The execution provider does not support device memory capacity queries"); + + uint64_t total{}; + uint64_t available{}; + OrtStatus* status = factory->GetDeviceMemoryCapacity(factory, ep_device, &total, &available); + if (status != nullptr) { + return status; + } + + *total_memory_bytes = total; + *available_memory_bytes = available; + return nullptr; + API_IMPL_END +} + namespace { OrtStatus* GetInputOutputMemoryInfo(const OrtSession* ort_session, InferenceSession::SessionInputOutputType type, @@ -4952,6 +4980,7 @@ static constexpr OrtApi ort_api_1_to_29 = { &OrtApis::SessionOptionsSetEpContextDataReadFunc, &OrtApis::GetModelPackageApi, + &OrtApis::EpDevice_GetMemoryCapacity, }; // OrtApiBase can never change as there is no way to know what version of OrtApiBase is returned by OrtGetApiBase. diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h index c448312884087..071d4c409dde7 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/onnxruntime/core/session/ort_apis.h @@ -755,6 +755,8 @@ ORT_API(const char*, GetRunConfigEntry, _In_ const OrtRunOptions* options, _In_z ORT_API(const OrtMemoryInfo*, EpDevice_MemoryInfo, _In_ const OrtEpDevice* ep_device, _In_ OrtDeviceMemoryType memory_type); +ORT_API_STATUS_IMPL(EpDevice_GetMemoryCapacity, _In_ const OrtEpDevice* ep_device, + _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes); ORT_API_STATUS_IMPL(CreateSharedAllocator, _In_ OrtEnv* env, _In_ const OrtEpDevice* ep_device, _In_ OrtDeviceMemoryType mem_type, _In_ OrtAllocatorType allocator_type, diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc b/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc index fb8a90f25ce13..86a6f6df95f19 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc @@ -32,6 +32,7 @@ EpFactoryInternal::EpFactoryInternal(std::unique_ptr impl OrtEpFactory::CreateDataTransfer = Forward::CreateDataTransfer; OrtEpFactory::IsStreamAware = Forward::IsStreamAware; OrtEpFactory::CreateSyncStreamForDevice = Forward::CreateSyncStreamForDevice; + OrtEpFactory::GetDeviceMemoryCapacity = Forward::GetDeviceMemoryCapacity; OrtEpFactory::CreateExternalResourceImporterForDevice = Forward::CreateExternalResourceImporterForDevice; OrtEpFactory::GetHardwareDeviceIncompatibilityDetails = Forward::GetHardwareDeviceIncompatibilityDetails; OrtEpFactory::InitGraphicsInterop = Forward::InitGraphicsInterop; diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal.h b/onnxruntime/core/session/plugin_ep/ep_factory_internal.h index cdfe5347185b7..5eca9efaa703d 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal.h @@ -80,6 +80,12 @@ class EpFactoryInternal : public OrtEpFactory { return impl_->CreateSyncStreamForDevice(memory_device, stream_options, stream); } + OrtStatus* GetDeviceMemoryCapacity(_In_ const OrtEpDevice* ep_device, + _Out_ uint64_t* total_memory_bytes, + _Out_ uint64_t* available_memory_bytes) noexcept { + return impl_->GetDeviceMemoryCapacity(ep_device, total_memory_bytes, available_memory_bytes); + } + OrtStatus* ValidateCompiledModelCompatibilityInfo(_In_reads_(num_devices) const OrtHardwareDevice* const* devices, _In_ size_t num_devices, _In_ const char* compatibility_info, diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h b/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h index cbadabfd8c3eb..59c7065e49d33 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h @@ -83,6 +83,14 @@ class EpFactoryInternalImpl { "CreateSyncStreamForDevice is not implemented for this EP factory."); } + virtual OrtStatus* GetDeviceMemoryCapacity( + _In_ const OrtEpDevice* /*ep_device*/, + _Out_ uint64_t* /*total_memory_bytes*/, + _Out_ uint64_t* /*available_memory_bytes*/) noexcept { + return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, + "GetDeviceMemoryCapacity is not implemented for this EP factory."); + } + virtual OrtStatus* CreateExternalResourceImporterForDevice( _In_ const OrtEpDevice* /*ep_device*/, _Outptr_result_maybenull_ OrtExternalResourceImporterImpl** importer) noexcept { diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h b/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h index add9f271b5ebc..78890d81bc0fd 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h @@ -65,6 +65,18 @@ class ProviderBridgeEpFactory : public EpFactoryInternalImpl { return ep_factory_.CreateSyncStreamForDevice(&ep_factory_, device, stream_options, stream); } + OrtStatus* GetDeviceMemoryCapacity(const OrtEpDevice* ep_device, + uint64_t* total_memory_bytes, + uint64_t* available_memory_bytes) noexcept override { + if (ep_factory_.ort_version_supported < 32 || + ep_factory_.GetDeviceMemoryCapacity == nullptr) { + return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, + "GetDeviceMemoryCapacity is not implemented for this EP factory."); + } + return ep_factory_.GetDeviceMemoryCapacity( + &ep_factory_, ep_device, total_memory_bytes, available_memory_bytes); + } + OrtStatus* CreateExternalResourceImporterForDevice( const OrtEpDevice* ep_device, OrtExternalResourceImporterImpl** importer) noexcept override { diff --git a/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h b/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h index 9dbf25a8251a4..0e4d174375b17 100644 --- a/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h +++ b/onnxruntime/core/session/plugin_ep/forward_to_factory_impl.h @@ -82,6 +82,15 @@ struct ForwardToFactoryImpl { return static_cast(this_ptr)->CreateSyncStreamForDevice(memory_device, stream_options, stream); } + static OrtStatus* ORT_API_CALL GetDeviceMemoryCapacity( + _In_ OrtEpFactory* this_ptr, + _In_ const OrtEpDevice* ep_device, + _Out_ uint64_t* total_memory_bytes, + _Out_ uint64_t* available_memory_bytes) noexcept { + return static_cast(this_ptr)->GetDeviceMemoryCapacity( + ep_device, total_memory_bytes, available_memory_bytes); + } + static OrtStatus* ORT_API_CALL CreateExternalResourceImporterForDevice( _In_ OrtEpFactory* this_ptr, _In_ const OrtEpDevice* ep_device, diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc index 2cf22da327a2a..3f63359bf3a7a 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc @@ -93,6 +93,7 @@ ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtL GetCustomOpDomains = GetCustomOpDomainsImpl; ValidateCompiledModelCompatibilityInfo = ValidateCompiledModelCompatibilityInfoImpl; SelectBestModelCandidate = SelectBestModelCandidateImpl; + GetDeviceMemoryCapacity = GetDeviceMemoryCapacityImpl; // setup the OrtMemoryInfo instances required by the EP. // We pretend the device the EP is running on is GPU. @@ -173,6 +174,15 @@ const char* ORT_API_CALL ExampleEpFactory::GetVersionImpl(const OrtEpFactory* th return factory->ep_version_.c_str(); } +/*static*/ +OrtStatus* ORT_API_CALL ExampleEpFactory::GetDeviceMemoryCapacityImpl( + OrtEpFactory* /*this_ptr*/, const OrtEpDevice* /*ep_device*/, + uint64_t* total_memory_bytes, uint64_t* available_memory_bytes) noexcept { + *total_memory_bytes = 16ULL * 1024 * 1024 * 1024; + *available_memory_bytes = 12ULL * 1024 * 1024 * 1024; + return nullptr; +} + /*static*/ OrtStatus* ORT_API_CALL ExampleEpFactory::GetSupportedDevicesImpl(OrtEpFactory* this_ptr, const OrtHardwareDevice* const* devices, diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h index 20103a7f45594..4c8bbbbe4373d 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h @@ -118,6 +118,12 @@ class ExampleEpFactory : public OrtEpFactory, public ApiPtrs { const OrtSessionOptions* session_options, size_t* selected_index) noexcept; + static OrtStatus* ORT_API_CALL GetDeviceMemoryCapacityImpl( + OrtEpFactory* this_ptr, + const OrtEpDevice* ep_device, + uint64_t* total_memory_bytes, + uint64_t* available_memory_bytes) noexcept; + const std::string ep_name_; // EP name const std::string vendor_{"Contoso"}; // EP vendor name const uint32_t vendor_id_{0xB357}; // EP vendor ID diff --git a/onnxruntime/test/autoep/test_registration.cc b/onnxruntime/test/autoep/test_registration.cc index f59ada9971637..c1ed2c223d9f7 100644 --- a/onnxruntime/test/autoep/test_registration.cc +++ b/onnxruntime/test/autoep/test_registration.cc @@ -7,6 +7,7 @@ #include #include "core/framework/data_types.h" +#include "core/providers/cuda/shared_inc/device_memory_capacity.h" #include "core/session/onnxruntime_cxx_api.h" #include "core/session/onnxruntime_ep_device_ep_metadata_keys.h" #include "core/session/onnxruntime_env_config_keys.h" @@ -34,6 +35,112 @@ extern "C" void ortenv_teardown(); namespace onnxruntime { namespace test { +namespace { +struct MemoryCapacityTestDriver { + using Device = int; + using Context = int; + + bool Step() noexcept { return ++step != failure_step; } + bool Init() noexcept { return Step(); } + bool DeviceGet(Device* device, int ordinal) noexcept { + queried_ordinal = ordinal; + *device = ordinal; + return Step(); + } + bool DeviceTotalMem(size_t* total, Device) noexcept { + *total = 16ULL * 1024 * 1024 * 1024; + return Step(); + } + bool PrimaryCtxRetain(Context* context, Device) noexcept { + *context = 42; + return retained = Step(); + } + bool CtxPushCurrent(Context context) noexcept { + if (!Step()) { + return false; + } + saved_context = current_context; + current_context = context; + return true; + } + bool MemGetInfo(size_t* available, size_t* total) noexcept { + queried_context = current_context; + *available = available_bytes; + *total = 14ULL * 1024 * 1024 * 1024; + return Step(); + } + bool CtxPopCurrent(Context* context) noexcept { + if (!Step()) { + return false; + } + *context = current_context; + current_context = saved_context; + return true; + } + bool PrimaryCtxRelease(Device) noexcept { + release_called = true; + if (!Step()) { + return false; + } + retained = false; + return true; + } + + int step = 0; + int failure_step = 0; + int current_context = 0; + int saved_context = 0; + int queried_context = 0; + int queried_ordinal = -1; + size_t available_bytes = 12ULL * 1024 * 1024 * 1024; + bool retained = false; + bool release_called = false; +}; +} // namespace + +TEST(OrtEpLibrary, MemoryCapacityPreservesCudaContextAndQueriesDeviceTotal) { + for (int original_context : {0, 73}) { + MemoryCapacityTestDriver driver; + driver.current_context = original_context; + uint64_t total = 0; + uint64_t available = 0; + ASSERT_EQ(cuda::QueryDeviceMemoryCapacity(driver, 3, total, available), nullptr); + EXPECT_EQ(driver.queried_ordinal, 3); + EXPECT_EQ(driver.queried_context, 42); + EXPECT_EQ(driver.current_context, original_context); + EXPECT_FALSE(driver.retained); + EXPECT_TRUE(driver.release_called); + EXPECT_EQ(total, 16ULL * 1024 * 1024 * 1024); + EXPECT_EQ(available, driver.available_bytes); + + driver.available_bytes /= 2; + ASSERT_EQ(cuda::QueryDeviceMemoryCapacity(driver, 3, total, available), nullptr); + EXPECT_EQ(available, driver.available_bytes); + EXPECT_EQ(driver.current_context, original_context); + } +} + +TEST(OrtEpLibrary, MemoryCapacityDriverFailuresPreserveOutputsAndCleanUp) { + for (int failure_step = 1; failure_step <= 8; ++failure_step) { + SCOPED_TRACE(failure_step); + MemoryCapacityTestDriver driver; + driver.failure_step = failure_step; + driver.current_context = 73; + uint64_t total = 7; + uint64_t available = 11; + ASSERT_NE(cuda::QueryDeviceMemoryCapacity(driver, 3, total, available), nullptr); + EXPECT_EQ(total, 7u); + EXPECT_EQ(available, 11u); + if (failure_step != 7) { + EXPECT_EQ(driver.current_context, 73); + } + if (failure_step != 8) { + EXPECT_FALSE(driver.retained); + } + EXPECT_EQ(driver.release_called, failure_step >= 5); + } +} + TEST(OrtEpLibrary, GetTensorDataTypeUsesOrtElementTypes) { const auto& ep_api = Ort::GetEpApi(); const struct { @@ -113,6 +220,16 @@ TEST(OrtEpLibrary, LoadUnloadPluginLibrary) { }); ASSERT_EQ(num_test_ep_devices, 1) << "Expected an OrtEpDevice to have been created by the test library."; + const OrtEpDevice* test_ep_device = *std::find_if( + ep_devices, ep_devices + num_devices, + [&](const OrtEpDevice* device) { return c_api->EpDevice_EpName(device) == ep_name; }); + uint64_t total_memory_bytes = 0; + uint64_t available_memory_bytes = 0; + ASSERT_ORTSTATUS_OK(c_api->EpDevice_GetMemoryCapacity( + test_ep_device, &total_memory_bytes, &available_memory_bytes)); + EXPECT_EQ(total_memory_bytes, 16ULL * 1024 * 1024 * 1024); + EXPECT_EQ(available_memory_bytes, 12ULL * 1024 * 1024 * 1024); + // and this should unload it ASSERT_ORTSTATUS_OK(Ort::GetApi().UnregisterExecutionProviderLibrary(*ort_env, registration_name.c_str())); @@ -135,6 +252,10 @@ TEST(OrtEpLibrary, LoadUnloadPluginLibraryCxxApi) { }); ASSERT_NE(test_ep_device, ep_devices.end()) << "Expected an OrtEpDevice to have been created by the test library."; + const auto [total_memory_bytes, available_memory_bytes] = test_ep_device->GetMemoryCapacity(); + EXPECT_EQ(total_memory_bytes, 16ULL * 1024 * 1024 * 1024); + EXPECT_EQ(available_memory_bytes, 12ULL * 1024 * 1024 * 1024); + // test all the C++ getters. // expected values are from \onnxruntime\test\autoep\library\example_plugin_ep\*.cc ASSERT_STREQ(test_ep_device->EpVendor(), "Contoso"); @@ -179,6 +300,31 @@ TEST(OrtEpLibrary, LoadUnloadPluginLibraryCxxApi) { ort_env->UnregisterExecutionProviderLibrary(registration_name.c_str()); } +TEST(OrtEpLibrary, UnsupportedMemoryCapacityQueryPreservesOutputs) { + const std::filesystem::path& library_path = Utils::example_ep_virt_gpu_info.library_path; + const std::string& registration_name = Utils::example_ep_virt_gpu_info.registration_name; + const std::string& ep_name = Utils::example_ep_virt_gpu_info.ep_name; + + ASSERT_ORTSTATUS_OK(Ort::GetApi().RegisterExecutionProviderLibrary( + *ort_env, registration_name.c_str(), library_path.c_str())); + + const auto ep_devices = ort_env->GetEpDevices(); + const auto test_ep_device = std::find_if( + ep_devices.begin(), ep_devices.end(), + [&](const Ort::ConstEpDevice& device) { return device.EpName() == ep_name; }); + ASSERT_NE(test_ep_device, ep_devices.end()); + + uint64_t total_memory_bytes = 7; + uint64_t available_memory_bytes = 11; + Ort::Status status{Ort::GetApi().EpDevice_GetMemoryCapacity( + *test_ep_device, &total_memory_bytes, &available_memory_bytes)}; + EXPECT_EQ(status.GetErrorCode(), ORT_NOT_IMPLEMENTED); + EXPECT_EQ(total_memory_bytes, 7u); + EXPECT_EQ(available_memory_bytes, 11u); + + ort_env->UnregisterExecutionProviderLibrary(registration_name.c_str()); +} + TEST(OrtEpLibrary, FailedRegistrationLeavesEnvironmentUnchanged) { const std::filesystem::path& library_path = Utils::example_ep_info.library_path; const std::string& registration_name = Utils::example_ep_info.registration_name; From 668a56522821af2002792bfe7e72f3841e52e25b Mon Sep 17 00:00:00 2001 From: Selena Yang <179177246+selenayang888@users.noreply.github.com> Date: Sun, 25 Oct 2026 23:30:47 -0700 Subject: [PATCH 2/4] Resolve the copilot comments --- cmake/onnxruntime_providers_cuda.cmake | 2 +- onnxruntime/test/autoep/test_registration.cc | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index c38138c94a770..02e88d9999727 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -385,7 +385,7 @@ add_dependencies(${target} onnxruntime_providers_shared ${onnxruntime_EXTERNAL_DEPENDENCIES}) if(onnxruntime_CUDA_MINIMAL) target_compile_definitions(${target} PRIVATE USE_CUDA_MINIMAL) - target_link_libraries(${target} PRIVATE ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED} Boost::mp11 safeint_interface CUDA::cudart) + target_link_libraries(${target} PRIVATE ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED} Boost::mp11 safeint_interface CUDA::cudart CUDA::cuda_driver) else() include(cudnn_frontend) # also defines CUDNN::* if (onnxruntime_USE_CUDA_NHWC_OPS) diff --git a/onnxruntime/test/autoep/test_registration.cc b/onnxruntime/test/autoep/test_registration.cc index c1ed2c223d9f7..3d2a4d9664df4 100644 --- a/onnxruntime/test/autoep/test_registration.cc +++ b/onnxruntime/test/autoep/test_registration.cc @@ -48,7 +48,7 @@ struct MemoryCapacityTestDriver { return Step(); } bool DeviceTotalMem(size_t* total, Device) noexcept { - *total = 16ULL * 1024 * 1024 * 1024; + *total = 16ULL * 1024 * 1024; return Step(); } bool PrimaryCtxRetain(Context* context, Device) noexcept { @@ -66,7 +66,7 @@ struct MemoryCapacityTestDriver { bool MemGetInfo(size_t* available, size_t* total) noexcept { queried_context = current_context; *available = available_bytes; - *total = 14ULL * 1024 * 1024 * 1024; + *total = 14ULL * 1024 * 1024; return Step(); } bool CtxPopCurrent(Context* context) noexcept { @@ -92,7 +92,7 @@ struct MemoryCapacityTestDriver { int saved_context = 0; int queried_context = 0; int queried_ordinal = -1; - size_t available_bytes = 12ULL * 1024 * 1024 * 1024; + size_t available_bytes = 12ULL * 1024 * 1024; bool retained = false; bool release_called = false; }; @@ -110,7 +110,7 @@ TEST(OrtEpLibrary, MemoryCapacityPreservesCudaContextAndQueriesDeviceTotal) { EXPECT_EQ(driver.current_context, original_context); EXPECT_FALSE(driver.retained); EXPECT_TRUE(driver.release_called); - EXPECT_EQ(total, 16ULL * 1024 * 1024 * 1024); + EXPECT_EQ(total, 16ULL * 1024 * 1024); EXPECT_EQ(available, driver.available_bytes); driver.available_bytes /= 2; From d9ba475a3082a615f897b7f2d35630fe03f50c09 Mon Sep 17 00:00:00 2001 From: Selena Yang <179177246+selenayang888@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:59:31 -0700 Subject: [PATCH 3/4] fix CI failure --- .github/workflows/android.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/android.yml b/.github/workflows/android.yml index 0bfc9178eff0a..11ba4dff6eda2 100644 --- a/.github/workflows/android.yml +++ b/.github/workflows/android.yml @@ -75,7 +75,7 @@ jobs: run: | set -e -x BINARY_SIZE_THRESHOLD_ARGS="" - BINARY_SIZE_THRESHOLD_IN_BYTES=1607680 + BINARY_SIZE_THRESHOLD_IN_BYTES=1608704 echo "Binary size threshold in bytes: ${BINARY_SIZE_THRESHOLD_IN_BYTES}" BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes ${BINARY_SIZE_THRESHOLD_IN_BYTES}" From 01734848e9047facdaf147b957bb310f2dd6bc44 Mon Sep 17 00:00:00 2001 From: Selena Yang <179177246+selenayang888@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:14:11 -0700 Subject: [PATCH 4/4] Move GetMemoryCapacity below CreateSyncStream in C++ wrappers Reorder the declaration and inline definition to address PR review feedback. --- .../onnxruntime/core/session/onnxruntime_cxx_api.h | 2 +- .../core/session/onnxruntime_cxx_inline.h | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 7917b8e1a2f0e..4f8bebac68453 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -1186,8 +1186,8 @@ struct EpDeviceImpl : Ort::detail::Base { ConstKeyValuePairs EpOptions() const; ConstHardwareDevice Device() const; ConstMemoryInfo GetMemoryInfo(OrtDeviceMemoryType memory_type) const; ///< Wraps EpDevice_MemoryInfo - std::pair GetMemoryCapacity() const; ///< Total and available device memory SyncStream CreateSyncStream(ConstKeyValuePairs stream_options = {}) const; /// Wraps EpDevice_CreateSyncStream + std::pair GetMemoryCapacity() const; ///< Total and available device memory }; } // namespace detail diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index e838355539297..0f62d9c6c42bb 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -756,6 +756,13 @@ inline ConstMemoryInfo EpDeviceImpl::GetMemoryInfo(OrtDeviceMemoryType memory return ConstMemoryInfo{mem_info}; } +template +inline SyncStream EpDeviceImpl::CreateSyncStream(ConstKeyValuePairs stream_options) const { + OrtSyncStream* stream = nullptr; + ThrowOnError(GetApi().CreateSyncStreamForEpDevice(this->p_, stream_options, &stream)); + return SyncStream{stream}; +} + template inline std::pair EpDeviceImpl::GetMemoryCapacity() const { uint64_t total_memory_bytes{}; @@ -763,13 +770,6 @@ inline std::pair EpDeviceImpl::GetMemoryCapacity() const ThrowOnError(GetApi().EpDevice_GetMemoryCapacity(this->p_, &total_memory_bytes, &available_memory_bytes)); return {total_memory_bytes, available_memory_bytes}; } - -template -inline SyncStream EpDeviceImpl::CreateSyncStream(ConstKeyValuePairs stream_options) const { - OrtSyncStream* stream = nullptr; - ThrowOnError(GetApi().CreateSyncStreamForEpDevice(this->p_, stream_options, &stream)); - return SyncStream{stream}; -} } // namespace detail inline EpDevice::EpDevice(OrtEpFactory& ep_factory, ConstHardwareDevice& hardware_device,