diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 34a15076c8a7f..6b1a3ed56b866 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -7517,6 +7517,32 @@ struct OrtApi { */ ORT_API2_STATUS(KernelContext_GetSyncStream, _In_ const OrtKernelContext* context, _Outptr_result_maybenull_ OrtSyncStream** out); + + /** \brief Set the source ONNX model as a byte buffer for weightless EPContext sessions. + * + * When creating a session from a weightless EPContext model, the EP may need access to the source model's + * initializer data. This function provides the source model as an in-memory byte buffer, for scenarios + * where the source model is not available as a file on disk (e.g., loaded from a package or downloaded). + * + * The caller retains ownership of the buffer and must ensure it remains valid for the lifetime of the session. + * + * \note If the source model is available as a file on disk, use the session config entry + * "ep.context_source_model_path" (kOrtSessionOptionEpContextSourceModelPath) instead. + * + * \note If both a buffer (via this function) and a file path (via "ep.context_source_model_path") are + * provided, the EP should prefer the buffer. The recommended EP precedence is: + * buffer > file path > "onnx_model_filename" EPContext node attribute. + * + * \param[in] options The OrtSessionOptions instance. + * \param[in] source_model_data Pointer to the source model byte buffer. + * \param[in] source_model_data_length Size of the byte buffer in bytes. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.29. + */ + ORT_API2_STATUS(SessionOptionsSetWeightlessSourceModelBuffer, _Inout_ OrtSessionOptions* options, + _In_ const void* source_model_data, _In_ size_t source_model_data_length); }; /* @@ -8362,6 +8388,31 @@ struct OrtCompileApi { ORT_API2_STATUS(ModelCompilationOptions_SetInputModel, _In_ OrtModelCompilationOptions* model_compile_options, _In_ const OrtModel* model); + + /** \brief Enable or disable weightless mode for model compilation. + * + * When enabled, the compiled EPContext model will not embed constant initializer data in the EP's + * compiled binary. Instead, the initializer data must be provided when creating a session from the + * compiled model, either from the source model (via the "onnx_model_filename" EPContext node attribute + * or the "ep.context_source_model_path" session option) or from externalized weights. + * + * This enables smaller compiled models and allows sharing initializer data across multiple compiled + * model variants (e.g., multi-platform caches for different hardware generations). + * + * ORT verifies that the target EP supports weightless mode during CompileModel() by calling + * OrtEp::GetWeightlessSupport(). If the EP does not support weightless mode, CompileModel() + * returns an error. + * + * \param[in] model_compile_options The OrtModelCompilationOptions instance. + * \param[in] use_weightless If true, enable weightless mode. If false, disable (default behavior). + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.29. + */ + ORT_API2_STATUS(ModelCompilationOptions_SetWeightlessEnabled, + _In_ OrtModelCompilationOptions* model_compile_options, + _In_ bool use_weightless); }; /** diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 55a4e36167e86..60b9bf0c3805c 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -1775,7 +1775,8 @@ struct ModelCompilationOptions : detail::Base { ModelCompilationOptions& SetGraphOptimizationLevel(GraphOptimizationLevel graph_optimization_level); ///< Wraps OrtApi::ModelCompilationOptions_SetGraphOptimizationLevel - ModelCompilationOptions& SetInputModel(const OrtModel* model); ///< Wraps OrtCompileApi::ModelCompilationOptions_SetInputModel + ModelCompilationOptions& SetInputModel(const OrtModel* model); ///< Wraps OrtCompileApi::ModelCompilationOptions_SetInputModel + ModelCompilationOptions& SetWeightlessEnabled(bool use_weightless); ///< Wraps OrtCompileApi::ModelCompilationOptions_SetWeightlessEnabled }; /** \brief Compiles an input model to generate a model with EPContext nodes that execute EP-specific kernels. Wraps OrtApi::CompileModels. diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index ed3abc0961be6..e8d5bf99735eb 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -1360,6 +1360,11 @@ inline ModelCompilationOptions& ModelCompilationOptions::SetInputModel(const Ort return *this; } +inline ModelCompilationOptions& ModelCompilationOptions::SetWeightlessEnabled(bool use_weightless) { + Ort::ThrowOnError(GetCompileApi().ModelCompilationOptions_SetWeightlessEnabled(this->p_, use_weightless)); + return *this; +} + namespace detail { template diff --git a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h index 683528304f7a0..0424a95c2d2c1 100644 --- a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h @@ -2077,6 +2077,32 @@ struct OrtEpApi { ORT_API2_STATUS(ProfilingEventsContainer_AddEvents, _In_ OrtProfilingEventsContainer* events_container, _In_reads_(num_events) const OrtProfilingEvent* const* events, _In_ size_t num_events); + + /** \brief Get the weightless source model byte buffer from session options. + * + * Returns the buffer and size set by SessionOptionsSetWeightlessSourceModelBuffer, or NULL/0 if not set. + * The EP can use this during CreateEp or Compile to access the source model for weightless + * EPContext model sessions. + * + * \note If the source model is provided as a file path, the EP should read the + * "ep.context_source_model_path" (kOrtSessionOptionEpContextSourceModelPath) session config entry + * via GetSessionConfigEntry instead. + * + * \note Recommended EP precedence for locating the source model: + * buffer (this API) > file path ("ep.context_source_model_path") > "onnx_model_filename" EPContext + * node attribute. + * + * \param[in] session_options The OrtSessionOptions instance. + * \param[out] source_model_data Output parameter set to the source model buffer, or NULL if not set. + * \param[out] source_model_data_length Output parameter set to the buffer size, or 0 if not set. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.29. + */ + ORT_API2_STATUS(SessionOptionsGetWeightlessSourceModelBuffer, _In_ const OrtSessionOptions* session_options, + _Outptr_result_maybenull_ const void** source_model_data, + _Out_ size_t* source_model_data_length); }; /** @@ -2111,6 +2137,26 @@ typedef enum OrtGraphCaptureNodeAssignmentPolicy { OrtGraphCaptureNodeAssignmentPolicy_ALLOW_CPU_FOR_SHAPES = 1, } OrtGraphCaptureNodeAssignmentPolicy; +/** + * \brief Describes the scope of an EP's weightless mode support. + * + * Returned by OrtEp::GetWeightlessSupport() to indicate which types of initializers + * the EP can operate on without copying. + * + * \since Version 1.29. + */ +typedef enum OrtWeightlessSupport { + /** EP does not support weightless mode. */ + OrtWeightlessSupport_NONE = 0, + + /** EP supports weightless mode for external initializers only. + * Internal initializers are still copied by the EP during compilation. */ + OrtWeightlessSupport_EXTERNAL_ONLY = 1, + + /** EP supports weightless mode for all initializers (internal and external). */ + OrtWeightlessSupport_ALL = 2, +} OrtWeightlessSupport; + /** * \brief The OrtEp struct provides functions to implement for an execution provider. * \since Version 1.22. @@ -2630,6 +2676,38 @@ struct OrtEp { * \since Version 1.27. */ ORT_API2_STATUS(ReleaseCapturedGraph, _In_ OrtEp* this_ptr, _In_ int graph_annotation_id); + + /** \brief Query the execution provider's weightless mode support. + * + * When weightless mode is enabled (via the "ep.enable_weightless" session option), ORT calls this function + * to determine the scope of the EP's weightless support. The EP returns an OrtWeightlessSupport value + * indicating whether it supports weightless mode for all initializers, external initializers only, or not + * at all. + * + * The EP's response may depend on the underlying hardware or driver capabilities. For example, an EP may + * support weightless mode for all initializers on newer hardware but only for external initializers on + * older hardware that requires weight transformation. + * + * EPs that support weightless mode should set drop_constant_initializers to false in OrtNodeFusionOptions + * so that ORT provides the initializer data as inputs to the compiled/fused node. The EP can then access + * these initializers at Compute() time via KernelContext_GetInput(). + * + * \note Extending the lifetime of initializer data obtained via ValueInfo_GetInitializerValue() during + * Compile() so that the EP can cache and reuse data pointers directly (without going through + * KernelContext) is planned but not yet implemented. Until then, KernelContext_GetInput() is the + * only supported way to access initializer data at Compute() time. + * + * \param[in] this_ptr The OrtEp instance. + * \param[out] support Output parameter set to the EP's weightless support scope. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \note Implementation of this function is optional. If set to NULL, ORT assumes the EP does not + * support weightless mode (equivalent to OrtWeightlessSupport_NONE). + * + * \since Version 1.29. + */ + ORT_API2_STATUS(GetWeightlessSupport, _In_ const OrtEp* this_ptr, _Out_ OrtWeightlessSupport* support); }; /** \brief The function signature that ORT will call to create OrtEpFactory instances. diff --git a/include/onnxruntime/core/session/onnxruntime_ep_device_ep_metadata_keys.h b/include/onnxruntime/core/session/onnxruntime_ep_device_ep_metadata_keys.h index cc83b7bca50c5..10e3627f37923 100644 --- a/include/onnxruntime/core/session/onnxruntime_ep_device_ep_metadata_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_ep_device_ep_metadata_keys.h @@ -30,3 +30,18 @@ static const char* const kOrtEpDevice_EpMetadataKey_LibraryPath = "library_path" // if this metadata key is not present. // - "1": OrtHardwareDevice is virtual. static const char* const kOrtHardwareDevice_MetadataKey_IsVirtual = "is_virtual"; + +// Key for the execution provider's weightless mode support on a specific device. +// Set by the EP during GetSupportedDevices() via CreateEpDevice() metadata. +// The app can read it via EpDevice_EpMetadata() to check device-specific weightless capability +// before calling ModelCompilationOptions_SetWeightlessEnabled(). +// +// Possible values: +// - "none": EP does not support weightless mode on this device. This is the assumed default value +// if this metadata key is not present. +// - "external_only": EP supports weightless mode for external initializers only (e.g., older +// hardware/driver that must transform internal constants). +// - "all": EP supports weightless mode for all initializers (internal and external). +// +// \since Version 1.29. +static const char* const kOrtEpDevice_EpMetadataKey_WeightlessSupport = "weightless_support"; diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index c29f43b5e0d30..7d01e305dc70f 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -583,8 +583,53 @@ static const char* const kOrtSessionOptionsRecordEpGraphAssignmentInfo = "sessio // Option values: // - "0": disable. (default) // - "1": enable. +// +// \deprecated Since version 1.29. Use "ep.enable_weightless" instead, which covers all initializers +// (internal and external) and works in both JIT and AOT flows. static const char* const kOrtSessionOptionEpEnableWeightlessEpContextNodes = "ep.enable_weightless_ep_context_nodes"; +// Enable weightless mode for all initializers (internal and external). +// +// When enabled, ONNX Runtime requests that the execution provider operate without embedding or copying +// constant initializers. +// +// This option works in both JIT (non-cached) and AOT (EPContext model) flows: +// - JIT: The EP should set drop_constant_initializers to false in OrtNodeFusionOptions so that ORT +// provides the initializer data as inputs to the compiled/fused node. The EP can then access these +// initializers at Compute() time via KernelContext_GetInput(). +// NOTE: Extending the lifetime of initializer data obtained via ValueInfo_GetInitializerValue() during +// Compile() so that the EP can cache and reuse data pointers directly is planned but not yet implemented. +// - AOT: ORT generates EPContext models with weightless EPContext nodes. The EP should use the +// "onnx_model_filename" EPContext node attribute or the "ep.context_source_model_path" session option +// to locate the source model's initializer data when creating a session from the compiled model. +// +// ORT checks that the EP supports weightless mode by calling OrtEpApi::GetWeightlessSupport(). +// If the EP does not support it, ORT returns an error. +// +// Option values: +// - "0": disable. (default) +// - "1": enable. +// +// \since Version 1.29. +static const char* const kOrtSessionOptionEpEnableWeightless = "ep.enable_weightless"; + +// Specifies the file path to the original (source) ONNX model when creating a session with a weightless +// EPContext model. +// +// When an EPContext model is generated with weightless mode ("ep.enable_weightless" = "1"), the compiled +// model may not contain the original initializer data. When creating a session from the compiled model, +// the EP needs to load the initializer data from the source model. This session option provides the +// runtime location of the source model, which may differ from the path used at compile time (stored in +// the EPContext node's "onnx_model_filename" attribute). +// +// If not set, the EP falls back to the "onnx_model_filename" attribute in the EPContext node. +// +// If the source model is available as a byte buffer rather than a file path, use +// OrtApi::SessionOptionsSetWeightlessSourceModelBuffer() instead. +// +// \since Version 1.29. +static const char* const kOrtSessionOptionEpContextSourceModelPath = "ep.context_source_model_path"; + // Controls the intra-op thread pool size for a session. // Value should be a base-10 int32 string. // Equivalent to OrtApi::SetIntraOpNumThreads. diff --git a/onnxruntime/core/session/abi_session_options_impl.h b/onnxruntime/core/session/abi_session_options_impl.h index c2a5d4d9985f8..0c1399e8a8e2b 100644 --- a/onnxruntime/core/session/abi_session_options_impl.h +++ b/onnxruntime/core/session/abi_session_options_impl.h @@ -39,4 +39,9 @@ struct OrtSessionOptions { // with GetProviderOptionPrefix returning 'ep.myep.' // CUDAExecutionProvider uses the stable short prefix 'ep.cuda.'. static std::string GetProviderOptionPrefix(const char* provider_name); + + // Weightless source model buffer for EPContext sessions. + // Set via SessionOptionsSetWeightlessSourceModelBuffer. + const void* weightless_source_model_data = nullptr; + size_t weightless_source_model_data_size = 0; }; diff --git a/onnxruntime/core/session/compile_api.cc b/onnxruntime/core/session/compile_api.cc index 54d26021d8c99..8e9bda27b9362 100644 --- a/onnxruntime/core/session/compile_api.cc +++ b/onnxruntime/core/session/compile_api.cc @@ -327,6 +327,22 @@ ORT_API_STATUS_IMPL(OrtCompileAPI::ModelCompilationOptions_SetInputModel, API_IMPL_END } +ORT_API_STATUS_IMPL(OrtCompileAPI::ModelCompilationOptions_SetWeightlessEnabled, + _In_ OrtModelCompilationOptions* ort_model_compile_options, + _In_ bool use_weightless) { + API_IMPL_BEGIN +#if !defined(ORT_MINIMAL_BUILD) + auto model_compile_options = reinterpret_cast(ort_model_compile_options); + ORT_API_RETURN_IF_STATUS_NOT_OK(model_compile_options->SetWeightlessEnabled(use_weightless)); + return nullptr; +#else + ORT_UNUSED_PARAMETER(ort_model_compile_options); + ORT_UNUSED_PARAMETER(use_weightless); + return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "Compile API is not supported in this build"); +#endif // !defined(ORT_MINIMAL_BUILD) + API_IMPL_END +} + ORT_API_STATUS_IMPL(OrtCompileAPI::CompileModel, _In_ const OrtEnv* env, _In_ const OrtModelCompilationOptions* ort_model_compile_options) { API_IMPL_BEGIN @@ -367,6 +383,8 @@ static constexpr OrtCompileApi ort_compile_api = { &OrtCompileAPI::ModelCompilationOptions_SetInputModel, // End of Version 24 - DO NOT MODIFY ABOVE + + &OrtCompileAPI::ModelCompilationOptions_SetWeightlessEnabled, }; // checks that we don't violate the rule that the functions must remain in the slots they were originally assigned diff --git a/onnxruntime/core/session/compile_api.h b/onnxruntime/core/session/compile_api.h index e8f171ee24295..627de23837215 100644 --- a/onnxruntime/core/session/compile_api.h +++ b/onnxruntime/core/session/compile_api.h @@ -45,4 +45,8 @@ ORT_API_STATUS_IMPL(ModelCompilationOptions_SetInputModel, _In_ OrtModelCompilationOptions* model_compile_options, _In_ const OrtModel* model); +ORT_API_STATUS_IMPL(ModelCompilationOptions_SetWeightlessEnabled, + _In_ OrtModelCompilationOptions* model_compile_options, + _In_ bool use_weightless); + } // namespace OrtCompileAPI diff --git a/onnxruntime/core/session/model_compilation_options.cc b/onnxruntime/core/session/model_compilation_options.cc index 9f6d1f9f1a9bc..524f1e67400d0 100644 --- a/onnxruntime/core/session/model_compilation_options.cc +++ b/onnxruntime/core/session/model_compilation_options.cc @@ -251,6 +251,13 @@ Status ModelCompilationOptions::SetGraphOptimizationLevel(GraphOptimizationLevel return Status::OK(); } +Status ModelCompilationOptions::SetWeightlessEnabled(bool use_weightless) { + ORT_RETURN_IF_ERROR( + session_options_.value.config_options.AddConfigEntry(kOrtSessionOptionEpEnableWeightless, + use_weightless ? "1" : "0")); + return Status::OK(); +} + Status ModelCompilationOptions::Check() const { const ConfigOptions& config_options = session_options_.value.config_options; diff --git a/onnxruntime/core/session/model_compilation_options.h b/onnxruntime/core/session/model_compilation_options.h index a15af565c4d54..1b0eb02595784 100644 --- a/onnxruntime/core/session/model_compilation_options.h +++ b/onnxruntime/core/session/model_compilation_options.h @@ -182,6 +182,14 @@ class ModelCompilationOptions { /// Status SetGraphOptimizationLevel(GraphOptimizationLevel graph_optimization_level); + /// + /// Enable weightless mode for model compilation. + /// When enabled, the compiled EPContext model will not embed constant initializer data. + /// + /// True to enable weightless mode + /// Status indicating potential error + Status SetWeightlessEnabled(bool use_weightless); + /// /// Checks if the compilation options described by this object are valid. /// diff --git a/onnxruntime/core/session/model_package/model_package_context.cc b/onnxruntime/core/session/model_package/model_package_context.cc index 351ef2498dcd5..6a3eecd514fae 100644 --- a/onnxruntime/core/session/model_package/model_package_context.cc +++ b/onnxruntime/core/session/model_package/model_package_context.cc @@ -33,7 +33,8 @@ bool IsModelPackagePathSessionOption(std::string_view key) { // Session-option config keys whose values are path references (sha256:, relative, or // absolute) that must be resolved against the model package. Add new path-valued keys here. return key == kOrtSessionOptionsModelExternalInitializersFileFolderPath || - key == kOrtSessionOptionEpContextFilePath; + key == kOrtSessionOptionEpContextFilePath || + key == kOrtSessionOptionEpContextSourceModelPath; } namespace { diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index e55af70915a39..915de32a813bf 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -2758,6 +2758,23 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsSetCustomJoinThreadFn, _Inout_ OrtSes API_IMPL_END } +ORT_API_STATUS_IMPL(OrtApis::SessionOptionsSetWeightlessSourceModelBuffer, _Inout_ OrtSessionOptions* options, + _In_ const void* source_model_data, _In_ size_t source_model_data_length) { + API_IMPL_BEGIN + if (source_model_data == nullptr) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Invalid source model: data pointer is null"); + } + + if (source_model_data_length == 0) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Invalid source model: data size is 0"); + } + + options->weightless_source_model_data = source_model_data; + options->weightless_source_model_data_size = source_model_data_length; + return nullptr; + API_IMPL_END +} + ORT_API(void, OrtApis::ReleaseValueInfo, _Frees_ptr_opt_ OrtValueInfo* value_info) { delete value_info; } @@ -4918,6 +4935,8 @@ static constexpr OrtApi ort_api_1_to_29 = { &OrtApis::GetExperimentalFunction, &OrtApis::KernelContext_GetSyncStream, // End of Version 28 - DO NOT MODIFY ABOVE (see above text for more information) + + &OrtApis::SessionOptionsSetWeightlessSourceModelBuffer, }; // 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 e747d0d0ab2d8..550c7d5c6a0f8 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/onnxruntime/core/session/ort_apis.h @@ -829,4 +829,8 @@ ORT_API_STATUS_IMPL(GetTensorElementTypeAndShapeDataReference, _In_ const OrtVal // Experimental API ORT_API(OrtExperimentalFnPtr, GetExperimentalFunction, _In_ const char* name); +// Weightless source model APIs +ORT_API_STATUS_IMPL(SessionOptionsSetWeightlessSourceModelBuffer, _Inout_ OrtSessionOptions* options, + _In_ const void* source_model_data, _In_ size_t source_model_data_length); + } // namespace OrtApis diff --git a/onnxruntime/core/session/plugin_ep/ep_api.cc b/onnxruntime/core/session/plugin_ep/ep_api.cc index d56f4299402b5..2f86e3a20b8d1 100644 --- a/onnxruntime/core/session/plugin_ep/ep_api.cc +++ b/onnxruntime/core/session/plugin_ep/ep_api.cc @@ -25,6 +25,7 @@ #include "core/session/abi_devices.h" #include "core/session/abi_ep_types.h" #include "core/session/abi_opschema.h" +#include "core/session/abi_session_options_impl.h" #include "core/session/environment.h" #include "core/session/onnxruntime_ep_device_ep_metadata_keys.h" #include "core/session/ort_apis.h" @@ -1198,6 +1199,16 @@ ORT_API_STATUS_IMPL(ProfilingEventsContainer_AddEvents, API_IMPL_END } +ORT_API_STATUS_IMPL(SessionOptionsGetWeightlessSourceModelBuffer, _In_ const OrtSessionOptions* session_options, + _Outptr_result_maybenull_ const void** source_model_data, + _Out_ size_t* source_model_data_length) { + API_IMPL_BEGIN + *source_model_data = session_options->weightless_source_model_data; + *source_model_data_length = session_options->weightless_source_model_data_size; + return nullptr; + API_IMPL_END +} + static constexpr OrtEpApi ort_ep_api = { // NOTE: ABI compatibility depends on the order within this struct so all additions must be at the end, // and no functions can be removed (the implementation needs to change to return an error). @@ -1287,6 +1298,8 @@ static constexpr OrtEpApi ort_ep_api = { &OrtExecutionProviderApi::ProfilingEvent_GetArgValue, &OrtExecutionProviderApi::ProfilingEventsContainer_AddEvents, // End of Version 25 - DO NOT MODIFY ABOVE + + &OrtExecutionProviderApi::SessionOptionsGetWeightlessSourceModelBuffer, }; // checks that we don't violate the rule that the functions must remain in the slots they were originally assigned diff --git a/onnxruntime/core/session/plugin_ep/ep_api.h b/onnxruntime/core/session/plugin_ep/ep_api.h index e32e267a75ba5..f1a95b7e0d8ca 100644 --- a/onnxruntime/core/session/plugin_ep/ep_api.h +++ b/onnxruntime/core/session/plugin_ep/ep_api.h @@ -179,4 +179,9 @@ ORT_API_STATUS_IMPL(ProfilingEvent_GetDurationUs, _In_ const OrtProfilingEvent* ORT_API_STATUS_IMPL(ProfilingEvent_GetArgValue, _In_ const OrtProfilingEvent* event, _In_ const char* key, _Outptr_result_maybenull_ const char** out); +// Weightless source model getter +ORT_API_STATUS_IMPL(SessionOptionsGetWeightlessSourceModelBuffer, _In_ const OrtSessionOptions* session_options, + _Outptr_result_maybenull_ const void** source_model_data, + _Out_ size_t* source_model_data_length); + } // namespace OrtExecutionProviderApi diff --git a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc index 3710901c42a5e..53556c2de34de 100644 --- a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc +++ b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc @@ -28,6 +28,7 @@ #include "core/session/plugin_ep/ep_kernel_registration.h" #include "core/session/plugin_ep/ep_event_profiling.h" #include "core/session/ort_apis.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "core/providers/partitioning_utils.h" namespace onnxruntime { @@ -185,6 +186,10 @@ PluginExecutionProvider::PluginExecutionProvider(UniqueOrtEp ep, const OrtSessio kernel_registry_(std::move(kernel_registry)) { generate_ep_ctx_model_ = session_options.value.GetEpContextGenerationOptions().enable; + // Record if the app requested weightless mode. Validation is deferred to Compile(). + weightless_requested_ = + session_options.value.config_options.GetConfigOrDefault(kOrtSessionOptionEpEnableWeightless, "0") != "0"; + // Extract EP-scoped session config entries. // Arena options go to session_arena_options_; the rest go to provider_options_. { @@ -582,6 +587,36 @@ Status PluginExecutionProvider::Compile(const std::vector& fu ORT_RETURN_IF(ort_ep_->ReleaseNodeComputeInfos == nullptr, "OrtEp for ", Type(), " did not provide a valid ReleaseNodeComputeInfos() function"); + // Validate EP weightless support if the app requested it. + if (weightless_requested_) { + if (ort_ep_->ort_version_supported >= 29) { + if (ort_ep_->GetWeightlessSupport == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "Weightless mode requested (ep.enable_weightless=1) but EP '", Type(), + "' does not implement GetWeightlessSupport."); + } + + OrtWeightlessSupport support = OrtWeightlessSupport_NONE; + auto* ort_status = ort_ep_->GetWeightlessSupport(ort_ep_.get(), &support); + if (ort_status != nullptr) { + return ToStatusAndRelease(ort_status); + } + + if (support == OrtWeightlessSupport_NONE) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "Weightless mode requested (ep.enable_weightless=1) but EP '", Type(), + "' does not support weightless mode on this device."); + } + } else { + LOGS(GetEpLoggerOrDefault(), INFO) << "Weightless mode requested (ep.enable_weightless=1) but EP '" + << Type() << "' was compiled with API version " + << ort_ep_->ort_version_supported + << " which predates GetWeightlessSupport (version 29). " + << "ORT cannot verify EP weightless support. " + << "The EP may still handle weightless via its own provider options."; + } + } + const logging::Logger& logger = GetEpLoggerOrDefault(); const size_t num_graphs = fused_nodes_and_graphs.size(); std::vector> api_graphs_holder; diff --git a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.h b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.h index 610f0b449c119..bd6fe2e109b8f 100644 --- a/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.h +++ b/onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.h @@ -174,6 +174,7 @@ class PluginExecutionProvider : public IExecutionProvider { std::vector ep_devices_; std::vector allocator_mem_infos_; bool generate_ep_ctx_model_ = false; + bool weightless_requested_ = false; // True if app set ep.enable_weightless=1 // Provider options extracted from session-level config (excluding arena.*). // Exposed through GetProviderOptions() so the framework reports the effective EP configuration. diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc b/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc index 90ad4e7976824..636172edf3ba7 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc @@ -190,6 +190,7 @@ ExampleEp::ExampleEp(ExampleEpFactory& factory, const std::string& name, const C GetCompiledModelCompatibilityInfo = GetCompiledModelCompatibilityInfoImpl; // compatibility info for compiled models Sync = SyncImpl; // optional. can be nullptr GetDefaultMemoryDevice = GetDefaultMemoryDeviceImpl; // optional. can be nullptr + GetWeightlessSupport = GetWeightlessSupportImpl; // weightless support IGNORE_ORTSTATUS(ort_api.Logger_LogMessage(&logger_, OrtLoggingLevel::ORT_LOGGING_LEVEL_INFO, @@ -203,6 +204,13 @@ const char* ORT_API_CALL ExampleEp ::GetNameImpl(const OrtEp* this_ptr) noexcept return ep->name_.c_str(); } +/*static*/ +OrtStatus* ORT_API_CALL ExampleEp::GetWeightlessSupportImpl(const OrtEp* /*this_ptr*/, + OrtWeightlessSupport* support) noexcept { + *support = OrtWeightlessSupport_ALL; + return nullptr; +} + bool ExampleEp::CopiesConstantInitializers() const { return !(config_.enable_ep_context && config_.enable_weightless_ep_context_nodes); } diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep.h b/onnxruntime/test/autoep/library/example_plugin_ep/ep.h index 4112abb723d39..94c2a5043a8d1 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep.h +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep.h @@ -83,6 +83,8 @@ class ExampleEp : public OrtEp, public ApiPtrs { private: static const char* ORT_API_CALL GetNameImpl(const OrtEp* this_ptr) noexcept; + static OrtStatus* ORT_API_CALL GetWeightlessSupportImpl(const OrtEp* this_ptr, + OrtWeightlessSupport* support) noexcept; static OrtStatus* ORT_API_CALL CreateAllocatorImpl(_In_ OrtEp* this_ptr, _In_ const OrtMemoryInfo* memory_info, 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 2fefaeffb5d34..7d4705d15183e 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc @@ -164,6 +164,8 @@ OrtStatus* ORT_API_CALL ExampleEpFactory::GetSupportedDevicesImpl(OrtEpFactory* // Example os_driver_version. A real EP would read the OS driver version from the device. // The format is a 4-part dot-separated version matching the DXCore DriverVersion property. factory->ort_api.AddKeyValuePair(ep_metadata, kOrtEpDevice_EpMetadataKey_OSDriverVersion, "31.0.101.1000"); + // Report weightless support for all initializers. + factory->ort_api.AddKeyValuePair(ep_metadata, kOrtEpDevice_EpMetadataKey_WeightlessSupport, "all"); factory->ort_api.AddKeyValuePair(ep_options, "run_really_fast", "true"); // OrtEpDevice copies ep_metadata and ep_options. diff --git a/onnxruntime/test/autoep/test_execution.cc b/onnxruntime/test/autoep/test_execution.cc index e95918c719324..f71a507e8f6ef 100644 --- a/onnxruntime/test/autoep/test_execution.cc +++ b/onnxruntime/test/autoep/test_execution.cc @@ -818,6 +818,107 @@ TEST(OrtEpLibrary, PluginEp_GenWeightlessEpContextModel) { } } +// Test weightless EP context model generation using the new ep.enable_weightless session option +// and ModelCompilationOptions_SetWeightlessEnabled API. +TEST(OrtEpLibrary, PluginEp_WeightlessAllInitializers_CompileApi) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + Ort::ConstEpDevice plugin_ep_device(example_ep.get()); + + { + const ORTCHAR_T* input_model_file = ORT_TSTR("testdata/mul_1.onnx"); + const ORTCHAR_T* output_model_file = ORT_TSTR("plugin_ep_weightless_all_init_ctx.onnx"); + std::filesystem::remove(output_model_file); + + std::unordered_map ep_options; + Ort::SessionOptions session_options; + + // Use the new unified weightless option (ep.enable_weightless) instead of the deprecated one. + session_options.AddConfigEntry(kOrtSessionOptionEpEnableWeightless, "1"); + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + // Create model compilation options and enable weightless cache via the CompileApi. + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetFlags(OrtCompileApiFlags_ERROR_IF_NO_NODES_COMPILED); + compile_options.SetInputModelPath(input_model_file); + compile_options.SetOutputModelPath(output_model_file); + compile_options.SetWeightlessEnabled(true); + + // Compile the model. + ASSERT_CXX_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + ASSERT_TRUE(std::filesystem::exists(output_model_file)); + + // Clean up. + std::filesystem::remove(output_model_file); + } +} + +// Test SessionOptionsSetWeightlessSourceModelBuffer with valid and invalid inputs. +TEST(OrtEpLibrary, PluginEp_WeightlessSourceModelBuffer_Validation) { + Ort::SessionOptions session_options; + const auto& api = Ort::GetApi(); + + // Valid buffer. + { + const char dummy_data[] = "dummy model bytes"; + OrtStatus* status = api.SessionOptionsSetWeightlessSourceModelBuffer( + session_options, dummy_data, sizeof(dummy_data)); + ASSERT_EQ(status, nullptr); + } + + // Null buffer should fail. + { + OrtStatus* status = api.SessionOptionsSetWeightlessSourceModelBuffer( + session_options, nullptr, 100); + ASSERT_NE(status, nullptr); + ASSERT_EQ(api.GetErrorCode(status), ORT_INVALID_ARGUMENT); + api.ReleaseStatus(status); + } + + // Zero-length buffer should fail. + { + const char dummy_data[] = "data"; + OrtStatus* status = api.SessionOptionsSetWeightlessSourceModelBuffer( + session_options, dummy_data, 0); + ASSERT_NE(status, nullptr); + ASSERT_EQ(api.GetErrorCode(status), ORT_INVALID_ARGUMENT); + api.ReleaseStatus(status); + } +} + +// Test that weightless mode returns an error when the EP does not implement GetWeightlessSupport. +// The virtual GPU EP is compiled with ORT_API_VERSION >= 29 but does not set GetWeightlessSupport, +// so the validation in Compile() should return EP_FAIL. +TEST(OrtEpLibrary, PluginEp_WeightlessMode_ErrorWhenEpDoesNotSupport) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_virt_gpu_info, example_ep)); + Ort::ConstEpDevice plugin_ep_device(example_ep.get()); + + const ORTCHAR_T* input_model_file = ORT_TSTR("testdata/add_mul_add.onnx"); + const ORTCHAR_T* output_model_file = ORT_TSTR("plugin_ep_weightless_error_test.onnx"); + std::filesystem::remove(output_model_file); + + std::unordered_map ep_options; + Ort::SessionOptions session_options; + + // Request weightless mode. + session_options.AddConfigEntry(kOrtSessionOptionEpEnableWeightless, "1"); + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetInputModelPath(input_model_file); + compile_options.SetOutputModelPath(output_model_file); + compile_options.SetWeightlessEnabled(true); + + // CompileModel should fail because the virtual GPU EP does not implement GetWeightlessSupport. + auto status = Ort::CompileModel(*ort_env, compile_options); + ASSERT_FALSE(status.IsOK()); + ASSERT_THAT(status.GetErrorMessage(), testing::HasSubstr("does not implement GetWeightlessSupport")); + + // Clean up. + std::filesystem::remove(output_model_file); +} + // Test loading a compiled model without registering the required EP with the session. // We expect to get an explicit error that says that an EPContext node generated by "example_ep" // was not assigned to the appropriate EP. diff --git a/onnxruntime/test/autoep/test_registration.cc b/onnxruntime/test/autoep/test_registration.cc index 40ac1670b07dc..158508cb18826 100644 --- a/onnxruntime/test/autoep/test_registration.cc +++ b/onnxruntime/test/autoep/test_registration.cc @@ -72,6 +72,8 @@ TEST(OrtEpLibrary, LoadUnloadPluginLibraryCxxApi) { ASSERT_STREQ(metadata.GetValue("supported_devices"), "CrackGriffin 7+"); // Verify the example plugin's expected os_driver_version value. ASSERT_STREQ(metadata.GetValue(kOrtEpDevice_EpMetadataKey_OSDriverVersion), "31.0.101.1000"); + // Verify the example plugin reports weightless support for all initializers. + ASSERT_STREQ(metadata.GetValue(kOrtEpDevice_EpMetadataKey_WeightlessSupport), "all"); auto options = test_ep_device->EpOptions(); ASSERT_STREQ(options.GetValue("run_really_fast"), "true");