From 4f01df4fa0574dbf80041cb19f11d095dc1f0d76 Mon Sep 17 00:00:00 2001 From: adrastogi Date: Wed, 25 Feb 2026 08:19:01 -0800 Subject: [PATCH 1/4] Add OrtModel input support for Compile API (#27332) ### Description This change adds a feature to the Compile API, allowing an in-memory OrtModel created via the Model Editor API to be compiled directly without first serializing to a file or buffer. ### Motivation and Context The Model Editor API and Compile API are both public C APIs in the ONNX Runtime, but until now there was no way to pass a programmatically constructed model directly to compilation. This change attempts to closes that gap (see #26750) and ensures the new code path behaves identically to the established file and buffer paths. --------- Co-authored-by: Aditya Rastogi Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- .../NativeCompileApiMethods.shared.cs | 12 + .../core/session/onnxruntime_c_api.h | 23 ++ .../core/session/onnxruntime_cxx_api.h | 2 + .../core/session/onnxruntime_cxx_inline.h | 5 + onnxruntime/core/session/compile_api.cc | 26 ++ onnxruntime/core/session/compile_api.h | 3 + .../core/session/model_compilation_options.cc | 76 +++- .../core/session/model_compilation_options.h | 26 +- onnxruntime/core/session/utils.cc | 91 +++++ .../test/shared_lib/test_model_builder_api.cc | 357 ++++++++++++++++++ 10 files changed, 612 insertions(+), 9 deletions(-) diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/NativeCompileApiMethods.shared.cs b/csharp/src/Microsoft.ML.OnnxRuntime/NativeCompileApiMethods.shared.cs index 84020d84c9e73..00ca25d0a6367 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/NativeCompileApiMethods.shared.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/NativeCompileApiMethods.shared.cs @@ -25,6 +25,7 @@ public struct OrtCompileApi public IntPtr ModelCompilationOptions_SetGraphOptimizationLevel; public IntPtr ModelCompilationOptions_SetOutputModelWriteFunc; public IntPtr ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc; + public IntPtr ModelCompilationOptions_SetInputModel; } internal class NativeMethods @@ -136,6 +137,12 @@ public DOrtModelCompilationOptions_SetOutputModelWriteFunc public DOrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc OrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc; + [UnmanagedFunctionPointer(CallingConvention.Winapi)] + public delegate IntPtr /* OrtStatus* */ DOrtModelCompilationOptions_SetInputModel( + IntPtr /* OrtModelCompilationOptions* */ options, + IntPtr /* const OrtModel* */ inputModel); + public DOrtModelCompilationOptions_SetInputModel OrtModelCompilationOptions_SetInputModel; + internal NativeMethods(OnnxRuntime.NativeMethods.DOrtGetCompileApi getCompileApi) { @@ -217,6 +224,11 @@ internal NativeMethods(OnnxRuntime.NativeMethods.DOrtGetCompileApi getCompileApi _compileApi.ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc, typeof(DOrtModelCompilationOptions_SetOutputModelGetInitializerLocationFunc)); + OrtModelCompilationOptions_SetInputModel = + (DOrtModelCompilationOptions_SetInputModel)Marshal.GetDelegateForFunctionPointer( + _compileApi.ModelCompilationOptions_SetInputModel, + typeof(DOrtModelCompilationOptions_SetInputModel)); + } } } diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 6ae1539d4c294..94fe2add50ad4 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -8015,6 +8015,29 @@ struct OrtCompileApi { ORT_API2_STATUS(ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc, _In_ OrtModelCompilationOptions* model_compile_options, _In_ OrtGetInitializerLocationFunc get_initializer_location_func, _In_ void* state); + + /** \brief Sets the OrtModel to compile. + * + * Sets an OrtModel created via the Model Editor API as the input for compilation. + * + * The input model's source (file path, memory buffer, or OrtModel) must be set with + * one of: ModelCompilationOptions_SetInputModelPath, ModelCompilationOptions_SetInputModelFromBuffer, + * or ModelCompilationOptions_SetInputModel. + * + * The OrtModel must have a complete graph with inputs, outputs, and nodes defined. + * The caller retains ownership of the OrtModel and must not release it until after + * CompileModel returns. + * + * \param[in] model_compile_options The OrtModelCompilationOptions instance. + * \param[in] model The OrtModel to compile. The model is borrowed (not copied or owned). + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.24. + */ + ORT_API2_STATUS(ModelCompilationOptions_SetInputModel, + _In_ OrtModelCompilationOptions* model_compile_options, + _In_ const OrtModel* model); }; /** diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 5cf8cf88bb054..14696969b1a5a 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -1599,6 +1599,8 @@ struct ModelCompilationOptions : detail::Base { ModelCompilationOptions& SetFlags(uint32_t flags); ///< Wraps OrtApi::ModelCompilationOptions_SetFlags ModelCompilationOptions& SetGraphOptimizationLevel(GraphOptimizationLevel graph_optimization_level); ///< Wraps OrtApi::ModelCompilationOptions_SetGraphOptimizationLevel + + ModelCompilationOptions& SetInputModel(const OrtModel* model); ///< Wraps OrtCompileApi::ModelCompilationOptions_SetInputModel }; /** \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 1a3e49130a1d1..c9fb5d5315b9c 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -1170,6 +1170,11 @@ inline ModelCompilationOptions& ModelCompilationOptions::SetGraphOptimizationLev return *this; } +inline ModelCompilationOptions& ModelCompilationOptions::SetInputModel(const OrtModel* model) { + Ort::ThrowOnError(GetCompileApi().ModelCompilationOptions_SetInputModel(this->p_, model)); + return *this; +} + namespace detail { template diff --git a/onnxruntime/core/session/compile_api.cc b/onnxruntime/core/session/compile_api.cc index 12127e9708255..54d26021d8c99 100644 --- a/onnxruntime/core/session/compile_api.cc +++ b/onnxruntime/core/session/compile_api.cc @@ -306,6 +306,27 @@ ORT_API_STATUS_IMPL(OrtCompileAPI::ModelCompilationOptions_SetGraphOptimizationL API_IMPL_END } +ORT_API_STATUS_IMPL(OrtCompileAPI::ModelCompilationOptions_SetInputModel, + _In_ OrtModelCompilationOptions* ort_model_compile_options, + _In_ const OrtModel* model) { + API_IMPL_BEGIN +#if !defined(ORT_MINIMAL_BUILD) + auto model_compile_options = reinterpret_cast(ort_model_compile_options); + + if (model == nullptr) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Invalid input model: OrtModel pointer is null"); + } + + model_compile_options->SetInputModel(model); + return nullptr; +#else + ORT_UNUSED_PARAMETER(ort_model_compile_options); + ORT_UNUSED_PARAMETER(model); + 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 @@ -343,6 +364,9 @@ static constexpr OrtCompileApi ort_compile_api = { &OrtCompileAPI::ModelCompilationOptions_SetOutputModelWriteFunc, &OrtCompileAPI::ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc, // End of Version 23 - DO NOT MODIFY ABOVE + + &OrtCompileAPI::ModelCompilationOptions_SetInputModel, + // End of Version 24 - DO NOT MODIFY ABOVE }; // checks that we don't violate the rule that the functions must remain in the slots they were originally assigned @@ -350,6 +374,8 @@ static_assert(offsetof(OrtCompileApi, CompileModel) / sizeof(void*) == 8, "Size of version 22 Api cannot change"); // initial version in ORT 1.22 static_assert(offsetof(OrtCompileApi, ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc) / sizeof(void*) == 13, "Size of version 23 of Api cannot change"); +static_assert(offsetof(OrtCompileApi, ModelCompilationOptions_SetInputModel) / sizeof(void*) == 14, + "Size of version 24 of Api cannot change"); ORT_API(const OrtCompileApi*, OrtCompileAPI::GetCompileApi) { return &ort_compile_api; diff --git a/onnxruntime/core/session/compile_api.h b/onnxruntime/core/session/compile_api.h index 34fa06340a7f9..e8f171ee24295 100644 --- a/onnxruntime/core/session/compile_api.h +++ b/onnxruntime/core/session/compile_api.h @@ -41,5 +41,8 @@ ORT_API_STATUS_IMPL(ModelCompilationOptions_SetOutputModelWriteFunc, ORT_API_STATUS_IMPL(ModelCompilationOptions_SetOutputModelGetInitializerLocationFunc, _In_ OrtModelCompilationOptions* model_compile_options, _In_ OrtGetInitializerLocationFunc get_initializer_location_func, _In_ void* state); +ORT_API_STATUS_IMPL(ModelCompilationOptions_SetInputModel, + _In_ OrtModelCompilationOptions* model_compile_options, + _In_ const OrtModel* model); } // namespace OrtCompileAPI diff --git a/onnxruntime/core/session/model_compilation_options.cc b/onnxruntime/core/session/model_compilation_options.cc index 468dacc30c054..efaf28fbeefc0 100644 --- a/onnxruntime/core/session/model_compilation_options.cc +++ b/onnxruntime/core/session/model_compilation_options.cc @@ -45,7 +45,16 @@ void ModelCompilationOptions::SetInputModelFromBuffer(const void* input_model_da input_model_data_size_ = input_model_data_size; } +void ModelCompilationOptions::SetInputModel(const OrtModel* model) { + ResetInputModelSettings(); + input_model_ = model; +} + Status ModelCompilationOptions::SetOutputModelPath(const std::filesystem::path& output_model_path) { + if (output_model_path.empty()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Output model path must not be empty."); + } + ConfigOptions& config_options = session_options_.value.config_options; epctx::ModelGenOptions& ep_context_gen_options = session_options_.value.ep_context_gen_options; @@ -186,10 +195,19 @@ size_t ModelCompilationOptions::GetInputModelDataSize() const { return input_model_data_size_; } +bool ModelCompilationOptions::InputModelComesFromOrtModel() const { + return input_model_ != nullptr; +} + +const OrtModel* ModelCompilationOptions::GetInputModel() const { + return input_model_; +} + void ModelCompilationOptions::ResetInputModelSettings() { input_model_path_.clear(); input_model_data_ = nullptr; input_model_data_size_ = 0; + input_model_ = nullptr; } Status ModelCompilationOptions::SetGraphOptimizationLevel(GraphOptimizationLevel graph_optimization_level) { @@ -229,16 +247,21 @@ Status ModelCompilationOptions::Check() const { // Check input model settings. const bool input_from_file = !input_model_path_.empty(); const bool input_from_memory = input_model_data_ != nullptr; + const bool input_from_model = input_model_ != nullptr; + + int input_source_count = (input_from_file ? 1 : 0) + + (input_from_memory ? 1 : 0) + + (input_from_model ? 1 : 0); - if (!input_from_file && !input_from_memory) { + if (input_source_count == 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input model to compile must be loaded from either a file or a memory buffer"); + "Input model to compile must be specified via file path, memory buffer, or OrtModel"); } - if (input_from_file && input_from_memory) { + if (input_source_count > 1) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input model to compile must be loaded from either a file or a memory buffer, ", - "but not both."); + "Input model to compile must be specified via exactly one of: ", + "file path, memory buffer, or OrtModel"); } if (input_from_file && !std::filesystem::exists(input_model_path_)) { @@ -249,12 +272,45 @@ Status ModelCompilationOptions::Check() const { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Buffer for input model data has size 0"); } + // Validate OrtModel input + if (input_from_model) { + if (input_model_->graph == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "OrtModel has no graph. Call AddGraphToModel before compilation."); + } + + if (input_model_->graph->GetNumInputs() == 0 || input_model_->graph->GetNumOutputs() == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "OrtModel graph must have at least one input and one output defined."); + } + + if (input_model_->domain_to_version.empty()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "OrtModel must specify at least one opset domain/version."); + } + + // Note: Additional validation (node connections, schema) happens during + // Model::LoadFromModelEditorApiModel -> Graph::Resolve() + } + // Check output model settings. const epctx::ModelGenOptions& ep_context_gen_options = session_options_.value.ep_context_gen_options; bool has_no_output_model_location = std::holds_alternative( ep_context_gen_options.output_model_location); - if (has_no_output_model_location && input_from_file) { + // Determine if we can derive an output path from the input + bool can_derive_output_path = input_from_file; + + // For OrtModel input, check if model_path is set in the graph using the virtual GetModelPath() method + // (avoids dynamic_cast which requires RTTI) + if (input_from_model && input_model_->graph) { + const ORTCHAR_T* model_path_cstr = input_model_->graph->GetModelPath(); + if (model_path_cstr && model_path_cstr[0] != ORT_TSTR('\0')) { + can_derive_output_path = true; + } + } + + if (has_no_output_model_location && can_derive_output_path) { // User did not specify an output file, an output buffer, or an output write function. We default to generating an // output file with a name based on the input file name, so do not return an error. return Status::OK(); @@ -294,7 +350,13 @@ Status ModelCompilationOptions::Check() const { } std::string ModelCompilationOptions::GetInputSourceForTelemetry() const { - return InputModelComesFromFile() ? "file" : "buffer"; + if (InputModelComesFromFile()) { + return "file"; + } + if (InputModelComesFromOrtModel()) { + return "ort_model"; + } + return "buffer"; } std::string ModelCompilationOptions::GetOutputTargetForTelemetry() const { diff --git a/onnxruntime/core/session/model_compilation_options.h b/onnxruntime/core/session/model_compilation_options.h index 4ba8712a6c9c7..47529e794677e 100644 --- a/onnxruntime/core/session/model_compilation_options.h +++ b/onnxruntime/core/session/model_compilation_options.h @@ -10,6 +10,7 @@ #include "core/common/status.h" #include "core/common/path_string.h" #include "core/framework/allocator.h" +#include "core/graph/model_editor_api_types.h" #include "core/session/abi_session_options_impl.h" #include "core/session/onnxruntime_c_api.h" #include "core/session/onnxruntime_session_options_config_keys.h" @@ -45,6 +46,14 @@ class ModelCompilationOptions { /// The size in bytes of the input model's buffer void SetInputModelFromBuffer(const void* input_model_data, size_t input_model_data_size); + /// + /// Sets the OrtModel to compile. + /// The OrtModel is borrowed (not copied) - caller must keep it alive until CompileModel returns. + /// Overrides any previous call to SetInputModelPath(), SetInputModelFromBuffer(), or SetInputModel(). + /// + /// The OrtModel to compile + void SetInputModel(const OrtModel* model); + /// /// Sets the file path to store the output/compiled ONNX model. /// Overrides any previous call to SetOutputModelPath() or SetOutputModelBuffer(). @@ -132,6 +141,18 @@ class ModelCompilationOptions { /// true if input model comes from a file bool InputModelComesFromFile() const; + /// + /// Returns true if the input model comes from an OrtModel pointer. + /// + /// true if input model comes from an OrtModel + bool InputModelComesFromOrtModel() const; + + /// + /// Returns the OrtModel to compile, or nullptr if not set. + /// + /// pointer to the OrtModel or nullptr + const OrtModel* GetInputModel() const; + /// /// Returns the buffer that contains the bytes for the input ONNX model. /// Returns nullptr if the input model is not stored in a buffer. @@ -162,9 +183,9 @@ class ModelCompilationOptions { // Telemetry helper methods /// - /// Returns a string describing the input source type: "file" or "buffer". + /// Returns a string describing the input source type: "file", "buffer", or "ort_model". /// - /// "file" or "buffer" + /// "file", "buffer", or "ort_model" std::string GetInputSourceForTelemetry() const; /// @@ -205,6 +226,7 @@ class ModelCompilationOptions { std::filesystem::path input_model_path_; const void* input_model_data_ = nullptr; size_t input_model_data_size_ = 0; + const OrtModel* input_model_ = nullptr; // Borrowed pointer }; } // namespace onnxruntime #endif // !defined(ORT_MINIMAL_BUILD) diff --git a/onnxruntime/core/session/utils.cc b/onnxruntime/core/session/utils.cc index a354cf26368d4..461d8e6a9d195 100644 --- a/onnxruntime/core/session/utils.cc +++ b/onnxruntime/core/session/utils.cc @@ -22,6 +22,7 @@ #include "core/session/onnxruntime_ep_device_ep_metadata_keys.h" #if !defined(ORT_MINIMAL_BUILD) +#include "core/graph/model_editor_api_types.h" #include "core/session/plugin_ep/ep_factory_internal.h" #include "core/session/plugin_ep/ep_plugin_provider_interfaces.h" #include "core/session/plugin_ep/ep_library_plugin.h" @@ -288,6 +289,90 @@ static OrtStatus* CreateSessionAndLoadModelImpl(_In_ const OrtSessionOptions* op return nullptr; } +#if !defined(ORT_MINIMAL_BUILD) +// Overload of CreateSessionAndLoadModelImpl that takes an OrtModel* directly. +// This ensures load-path parity with file/buffer inputs by running the same checks +// (ORT_LOAD_CONFIG_FROM_MODEL, EP-context output validation, custom domain wiring). +static OrtStatus* CreateSessionAndLoadModelImpl(_In_ const OrtSessionOptions* options, + const onnxruntime::Environment& env, + _In_ const OrtModel* model, + std::unique_ptr& sess) { + if (model == nullptr) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "OrtModel pointer is null"); + } + + // Check EPContext model generation options - OrtModel has no file path by default, + // so we need explicit output location or embedded model path. + if (options) { + epctx::ModelGenOptions ep_ctx_gen_options = options->value.GetEpContextGenerationOptions(); + + if (ep_ctx_gen_options.enable) { + auto* output_model_path = ep_ctx_gen_options.TryGetOutputModelPath(); + + // Check if OrtModel has a model_path set + bool has_model_path = false; + if (model->graph) { + const ORTCHAR_T* model_path_cstr = model->graph->GetModelPath(); + has_model_path = model_path_cstr && model_path_cstr[0] != ORT_TSTR('\0'); + } + + // If there's no model path and no output location, fail early + if (!has_model_path && + (!ep_ctx_gen_options.HasOutputModelLocation() || + (output_model_path != nullptr && output_model_path->empty()))) { + return OrtApis::CreateStatus(ORT_FAIL, + "OrtModel has no model_path set and no valid output location was specified " + "for EPContext model generation. " + "SetOutputModelPath/SetOutputModelBuffer, or set the model_path on the " + "OrtGraph before adding it to OrtModel."); + } + } + } + + sess = std::make_unique( + options == nullptr ? onnxruntime::SessionOptions() : options->value, + env); + +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) + // Add custom domains + if (options && !options->custom_op_domains_.empty()) { + ORT_API_RETURN_IF_STATUS_NOT_OK(sess->AddCustomOpDomains(options->custom_op_domains_)); + } +#endif + +#if !defined(ORT_MINIMAL_BUILD) + // Add custom domains for all OrtEpDevice instances to inference session. + // The custom domains should be registered before model load for ORT to validate the custom ops. + // This mirrors the same block in the file/buffer overload to maintain load-path parity. + if (options != nullptr && + options->provider_factories.empty() && + options->value.ep_selection_policy.enable) { + InlinedVector all_ep_custom_op_domains; + + for (const OrtEpDevice* ep_device : env.GetOrtEpDevices()) { + InlinedVector domains; + ORT_API_RETURN_IF_STATUS_NOT_OK(GetCustomOpDomainsFromEpDevice(*ep_device, domains)); + + for (auto domain : domains) { + if (ShouldAddDomain(domain, options->custom_op_domains_)) { + all_ep_custom_op_domains.push_back(domain); + } + } + } + + if (!all_ep_custom_op_domains.empty()) { + ORT_API_RETURN_IF_STATUS_NOT_OK(sess->AddCustomOpDomains(all_ep_custom_op_domains)); + } + } +#endif // !defined(ORT_MINIMAL_BUILD) + + // Load from OrtModel + ORT_API_RETURN_IF_STATUS_NOT_OK(sess->Load(*model)); + + return nullptr; +} +#endif // !defined(ORT_MINIMAL_BUILD) + // Creates an InferenceSession and loads the model. // Caller should provide either model_path, or modal_data + model_data_length. OrtStatus* CreateSessionAndLoadModel(_In_ const OrtSessionOptions* options, @@ -491,6 +576,12 @@ Status CompileModel(const Environment& env, const ModelCompilationOptions& model status = ToStatusAndRelease(CreateSessionAndLoadModelImpl(session_options, env, input_model_path.c_str(), nullptr, 0, session)); + } else if (model_compile_options.InputModelComesFromOrtModel()) { + // Use the OrtModel overload of CreateSessionAndLoadModelImpl to maintain load-path parity + // with file/buffer inputs (same checks for ORT_LOAD_CONFIG_FROM_MODEL, EP-context output, etc.) + const OrtModel* input_model = model_compile_options.GetInputModel(); + status = ToStatusAndRelease(CreateSessionAndLoadModelImpl(session_options, env, + input_model, session)); } else { status = ToStatusAndRelease(CreateSessionAndLoadModelImpl(session_options, env, nullptr, model_compile_options.GetInputModelData(), diff --git a/onnxruntime/test/shared_lib/test_model_builder_api.cc b/onnxruntime/test/shared_lib/test_model_builder_api.cc index 018204bd1dfb0..ea5e889ad67a4 100644 --- a/onnxruntime/test/shared_lib/test_model_builder_api.cc +++ b/onnxruntime/test/shared_lib/test_model_builder_api.cc @@ -18,6 +18,7 @@ #include "test/shared_lib/test_fixture.h" #include "test/shared_lib/utils.h" +#include "test/util/include/api_asserts.h" #include "test/util/include/test_allocator.h" #include "onnxruntime_config.h" // generated file in build output dir @@ -725,3 +726,359 @@ TEST(ModelEditorAPITest, CreateTypeInfo) { api.ReleaseTypeInfo(base_tensor_type_info); } + +// +// Tests for Model Editor API + Compile API integration +// + +namespace { +// Helper to create a simple model for testing with Model Editor API +// Creates a model with a Gemm operation: Z = X * Y where X is input and Y is initializer +Ort::Model CreateSimpleGemmModel(std::vector>>& weights) { + Ort::Graph graph; + + std::vector graph_inputs; + std::vector graph_outputs; + + // Input: X is 3x4 + std::vector input_dims({3, 4}); + TensorTypeAndShapeInfo input_tensor_info(ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, + input_dims); + auto input_type_info = TypeInfo::CreateTensorInfo(input_tensor_info.GetConst()); + graph_inputs.emplace_back("X", input_type_info.GetConst()); + + // Output: Z is 3x8 + std::vector output_dims = {3, 8}; + TensorTypeAndShapeInfo output_tensor_info(ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, + output_dims); + auto output_type_info = TypeInfo::CreateTensorInfo(output_tensor_info.GetConst()); + graph_outputs.emplace_back("Z", output_type_info.GetConst()); + + graph.SetInputs(graph_inputs); + graph.SetOutputs(graph_outputs); + + // Gemm node with alpha=2.0 + std::vector attributes; + float alpha_value = 2.0; + attributes.push_back(OpAttr("alpha", &alpha_value, 1, OrtOpAttrType::ORT_OP_ATTR_FLOAT)); + + Node node("Gemm", onnxruntime::kOnnxDomain, "Gemm1", {"X", "Y"}, {"Z"}, attributes); + graph.AddNode(node); + + // Y initializer: 4x8 + std::vector y_dims = {4, 8}; + weights.emplace_back(std::make_unique>(32)); + auto& y_values = *weights.back(); + std::iota(y_values.begin(), y_values.end(), 1.0f); + + auto info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto y_tensor = Value::CreateTensor(info, y_values.data(), y_values.size(), y_dims.data(), y_dims.size()); + graph.AddInitializer("Y", y_tensor, /*data is external*/ true); + + std::vector opsets{{onnxruntime::kOnnxDomain, 18}}; + Model model(opsets); + model.AddGraph(graph); + + return model; +} + +// Helper to run inference on the simple Gemm model and verify all output values. +// Model is Z = 2.0 * X * Y where X is 3x4 (all ones) and Y is 4x8 (iota 1..32). +// Expected output: each row is 2 * column_sums_of_Y = {104, 112, 120, 128, 136, 144, 152, 160}. +void RunAndVerifySimpleGemmModel(const Ort::Model& model) { + Ort::SessionOptions session_options; + Ort::Session session(*ort_env, model, session_options); + ASSERT_EQ(session.GetInputCount(), 1u); + ASSERT_EQ(session.GetOutputCount(), 1u); + + std::vector input_data(3 * 4, 1.0f); + std::vector input_dims = {3, 4}; + auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto input_tensor = Ort::Value::CreateTensor(memory_info, input_data.data(), input_data.size(), + input_dims.data(), input_dims.size()); + + const char* input_names[] = {"X"}; + const char* output_names[] = {"Z"}; + auto outputs = session.Run(Ort::RunOptions{}, input_names, &input_tensor, 1, output_names, 1); + ASSERT_EQ(outputs.size(), 1u); + ASSERT_TRUE(outputs[0].IsTensor()); + + auto output_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape(); + ASSERT_EQ(output_shape, (std::vector{3, 8})); + + const float* output_data = outputs[0].GetTensorData(); + // alpha=2.0, X is all ones, so each output row = 2 * sum of each column of Y (iota 1..32 in 4x8) + const std::vector expected_row = {104.0f, 112.0f, 120.0f, 128.0f, 136.0f, 144.0f, 152.0f, 160.0f}; + for (int row = 0; row < 3; ++row) { + for (int col = 0; col < 8; ++col) { + EXPECT_FLOAT_EQ(output_data[row * 8 + col], expected_row[col]) + << "Mismatch at row=" << row << " col=" << col; + } + } +} +} // namespace + +// Test basic compilation from OrtModel +TEST(ModelEditorCompileAPITest, BasicCompileFromOrtModel) { + std::vector>> weights; + auto model = CreateSimpleGemmModel(weights); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + + // Set the OrtModel as input + compile_options.SetInputModel(static_cast(model)); + + // Set output to buffer - use embed mode for simplicity + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + // Compile should succeed (note: may not produce EPContext nodes without specific EP, but validation passes) + ASSERT_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + + // Verify output was produced + EXPECT_NE(output_buffer, nullptr); + EXPECT_GT(output_size, 0u); + + // Cleanup + if (output_buffer != nullptr) { + allocator->Free(output_buffer); + } + + // Verify the model still produces correct inference results after compilation + RunAndVerifySimpleGemmModel(model); +} +TEST(ModelEditorCompileAPITest, CompileFromNullModel_Fails) { + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + + try { + compile_options.SetInputModel(nullptr); + FAIL() << "Expected exception for null model pointer"; + } catch (const Ort::Exception& e) { + EXPECT_THAT(e.what(), ::testing::HasSubstr("null")); + } +} + +// Test validation: model with no graph +TEST(ModelEditorCompileAPITest, CompileFromModelWithNoGraph_Fails) { + // Create a model but don't add a graph + std::vector opsets{{onnxruntime::kOnnxDomain, 18}}; + Model model(opsets); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + + compile_options.SetInputModel(static_cast(model)); + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + Ort::Status status = Ort::CompileModel(*ort_env, compile_options); + EXPECT_FALSE(status.IsOK()) << "Expected CompileModel to fail for model with no graph"; + EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("graph")); +} + +// Test validation: model with empty inputs/outputs +TEST(ModelEditorCompileAPITest, CompileFromModelWithEmptyInputsOutputs_Fails) { + // Create a model with a graph that has no inputs or outputs + Ort::Graph graph; + // Don't set inputs or outputs + + std::vector opsets{{onnxruntime::kOnnxDomain, 18}}; + Model model(opsets); + model.AddGraph(graph); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + + compile_options.SetInputModel(static_cast(model)); + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + Ort::Status status = Ort::CompileModel(*ort_env, compile_options); + EXPECT_FALSE(status.IsOK()) << "Expected CompileModel to fail for model with empty inputs/outputs"; + EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("input")); +} + +// Test: model can be reused after compilation. +// NOTE: This is not an explicit API guarantee. It documents current behavior so that if a future change +// breaks model reuse, the regression is surfaced and can be evaluated. +TEST(ModelEditorCompileAPITest, ModelCanBeReusedAfterCompilation) { + std::vector>> weights; + auto model = CreateSimpleGemmModel(weights); + + // First compilation + { + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetInputModel(static_cast(model)); + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + ASSERT_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + + if (output_buffer != nullptr) { + allocator->Free(output_buffer); + } + } + + // Second compilation with same model + { + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetInputModel(static_cast(model)); + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + ASSERT_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + + if (output_buffer != nullptr) { + allocator->Free(output_buffer); + } + } + + // Model should still be usable for creating a session and running inference + RunAndVerifySimpleGemmModel(model); +} + +// Test: SetInputModel overrides previous input source (file path) +TEST(ModelEditorCompileAPITest, SetInputModelOverridesPreviousInputPath) { + std::vector>> weights; + auto model = CreateSimpleGemmModel(weights); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + + // First set a file path (doesn't need to exist since we'll override it) + compile_options.SetInputModelPath(ORT_TSTR("nonexistent_file.onnx")); + + // Then override with OrtModel + compile_options.SetInputModel(static_cast(model)); + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + // Should use the OrtModel, not the nonexistent file + ASSERT_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + + if (output_buffer != nullptr) { + allocator->Free(output_buffer); + } +} + +// Test: SetInputModelPath overrides previous OrtModel setting +TEST(ModelEditorCompileAPITest, SetInputModelPathOverridesPreviousModel) { + std::vector>> weights; + auto model = CreateSimpleGemmModel(weights); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + + // First set an OrtModel + compile_options.SetInputModel(static_cast(model)); + + // Then override with a real file path + compile_options.SetInputModelPath(ORT_TSTR("testdata/matmul_1.onnx")); + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + // Should use the file path, not the OrtModel + ASSERT_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + + if (output_buffer != nullptr) { + allocator->Free(output_buffer); + } +} + +// Test: Compile with output to file +TEST(ModelEditorCompileAPITest, CompileFromOrtModelToFile) { + std::vector>> weights; + auto model = CreateSimpleGemmModel(weights); + + auto output_path = ORT_TSTR("test_compile_from_ortmodel_output.onnx"); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + + compile_options.SetInputModel(static_cast(model)); + compile_options.SetOutputModelPath(output_path); + compile_options.SetEpContextEmbedMode(true); + + ASSERT_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + + // Verify output file exists + EXPECT_TRUE(std::filesystem::exists(output_path)); + + // Verify the output model can be loaded + Ort::Session session(*ort_env, output_path, Ort::SessionOptions()); + EXPECT_GE(session.GetInputCount(), 1u); + EXPECT_GE(session.GetOutputCount(), 1u); + + // Cleanup + std::filesystem::remove(output_path); +} + +// Test: Validation error for OrtModel with no model_path, no output location, and no embed mode. +TEST(ModelEditorCompileAPITest, NoOutputLocationNoModelPathFails) { + std::vector>> weights; + auto model = CreateSimpleGemmModel(weights); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetInputModel(static_cast(model)); + // Intentionally do NOT call SetEpContextEmbedMode, SetOutputModelPath, or SetOutputModelBuffer + + Ort::Status status = Ort::CompileModel(*ort_env, compile_options); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("output")); +} + +// Test: Setting embed mode with buffer output satisfies the output location requirement +// for OrtModel with no model_path. +TEST(ModelEditorCompileAPITest, EmbedModeWithBufferOutputSatisfiesValidation) { + std::vector>> weights; + auto model = CreateSimpleGemmModel(weights); + + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetInputModel(static_cast(model)); + compile_options.SetEpContextEmbedMode(true); + + std::unique_ptr allocator = std::make_unique(); + void* output_buffer = nullptr; + size_t output_size = 0; + compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size); + + ASSERT_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + + if (output_buffer != nullptr) { + allocator->Free(output_buffer); + } +} From 6d4564c1ed5d0b599dcf0dffc39d570b2c6ac68b Mon Sep 17 00:00:00 2001 From: adrastogi Date: Sun, 29 Mar 2026 20:02:57 -0700 Subject: [PATCH 2/4] Fix overflow in DmlGraphFusionHelper::ProcessInputData (#27815) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This change tries to address a problem in the DML EP where AlignToPow2 rounded up tensorByteSize to a 4-byte boundary before the data was read from the source buffer. This caused CreateCpuResource, CreateResource, WriteToFile, and the inputRawData vector construction to read 1–3 bytes past the end of the original tensor data. CreateResource and CreateCpuResource already independently align the D3D12 resource descriptor size, so they work correctly with the original (unaligned) byte count. The fix is to move the alignment to the location where it's needed. This is required because it addresses a crash / incorrect behavior in the DML EP. --- .../src/DmlGraphFusionHelper.cpp | 7 ++-- .../cpu/tensor/quantize_linear_test.cc | 41 +++++++++++++++++++ 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp index 6bd7de0fba5cb..4ddf8b8640376 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp @@ -232,8 +232,6 @@ namespace DmlGraphFusionHelper } } - // Tensor sizes in DML must be a multiple of 4 bytes large. - tensorByteSize = AlignToPow2(tensorByteSize, 4); if(graphSerializationEnabled) { WriteToFile(modelName, ConvertToWString(iter->first) + L".bin", reinterpret_cast(tensorPtr), tensorByteSize); @@ -264,9 +262,10 @@ namespace DmlGraphFusionHelper initializeInputBuffer = CreateCpuResource(providerImpl, tensorPtr, tensorByteSize); } - // Set the binding for operator initialization to the buffer + // Set the binding for operator initialization to the buffer. + // DML requires buffer binding sizes to be a multiple of 4 bytes. initInputBindings[i].Buffer = initializeInputBuffer.Get(); - initInputBindings[i].SizeInBytes = tensorByteSize; + initInputBindings[i].SizeInBytes = AlignToPow2(tensorByteSize, 4); initializeResourceRefs.push_back(std::move(initializeInputBuffer)); } diff --git a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc index bf632d0b3bc40..62fe0ffbc0581 100644 --- a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc @@ -60,6 +60,47 @@ TEST(DequantizeLinearOpTest, Int8_Large) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); } +TEST(DequantizeLinearOpTest, Int4_LargeInitializerInput) { + OpTester test("DequantizeLinear", 21); + std::vector dims{1024}; + + std::vector x_vals(Int4x2::CalcNumInt4Pairs(static_cast(dims[0])), Int4x2{}); + std::vector expected_y_vals(static_cast(dims[0]), 0.f); + + test.AddInput("x", dims, x_vals, true); + test.AddInput("x_scale", {}, {1.0f}); + test.AddInput("x_zero_point", {}, {Int4x2(0, 0)}); + test.AddOutput("y", dims, expected_y_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +// Regression test: int8 tensor whose byte size is not a multiple of 4. +// DML graph fusion rounds tensor sizes to a multiple of 4 via AlignToPow2. +// If the original buffer is not padded, the subsequent memcpy reads past the +// allocation boundary (heap-buffer-overflow detectable with ASan). +// Mirrors the WebNN PoC: dequantizeLinear with int8[135] (135 % 4 != 0). +TEST(DequantizeLinearOpTest, Int8_NonAlignedSize_Initializer) { + OpTester test("DequantizeLinear", 10); + constexpr int64_t kNumElements = 135; // 135 bytes, AlignToPow2(135,4)=136 + + std::vector x_data(kNumElements); + std::vector y_expected(kNumElements); + const float scale = 0.5f; + const int8_t zero_point = 0; + for (int64_t i = 0; i < kNumElements; ++i) { + x_data[i] = static_cast(i % 127); + y_expected[i] = (x_data[i] - zero_point) * scale; + } + + // Mark all inputs as initializers so they go through DML's ProcessInputData + // → UnpackInitializer → AlignToPow2 code path during graph fusion. + test.AddInput("x", {kNumElements}, x_data, /*is_initializer=*/true); + test.AddInput("x_scale", {1}, {scale}, /*is_initializer=*/true); + test.AddInput("x_zero_point", {1}, {zero_point}, /*is_initializer=*/true); + test.AddOutput("y", {kNumElements}, y_expected); + + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} // scalar zero & scale with int4 TEST(DequantizeLinearOpTest, Int4) { OpTester test("DequantizeLinear", 21); From bd5d980aebb27d807e6d26a106c30cfcca6c3978 Mon Sep 17 00:00:00 2001 From: adrastogi Date: Sun, 29 Mar 2026 17:29:51 -0700 Subject: [PATCH 3/4] Fix new-delete mismatch in DML EP's QuantizeLinear operator (#27823) ### Description DmlOperatorQuantization21 was missing the tensor reshaping logic that the older DmlOperatorElementwiseQLinear already had. Scalar scale tensors get padded to 4D, but a 5D input stays 5D. DML rejects the dimension mismatch with E_INVALIDARG, and the resulting exception unwind triggers a sized-delete bug in WRL's MakeAllocator which address sanitizer detects. The fix is to port the same logic from the DmlOperatorElementwiseQLinear into this path, so that the dimensions match. ### Motivation and Context This is required to ensure the DML EP correctly handles this scenario. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/AbiCustomRegistry.cpp | 2 +- .../src/BucketizedBufferAllocator.cpp | 2 +- .../src/DmlCommittedResourceAllocator.cpp | 2 +- .../src/DmlExternalBufferAllocator.h | 4 +- .../src/ExecutionProvider.cpp | 8 +- .../src/GraphDescBuilder.cpp | 2 +- .../src/MLOperatorAuthorImpl.cpp | 32 ++-- .../src/Operators/DmlDFT.h | 6 +- .../src/Operators/DmlGridSample.h | 6 +- .../src/Operators/DmlOperator.cpp | 67 +++++++++ .../src/Operators/DmlOperator.h | 9 ++ .../src/Operators/DmlOperatorElementWise.cpp | 61 +------- .../src/Operators/DmlOperatorNonZero.cpp | 6 +- .../src/Operators/DmlSTFT.h | 8 +- .../src/Operators/OperatorRegistration.cpp | 6 +- .../src/SafeMakeOrThrow.h | 37 +++++ .../dml/DmlExecutionProvider/src/precomp.h | 1 + .../MLOperatorAuthorHelper.h | 3 +- .../SchemaInferenceOverrider.h | 3 +- .../providers/dml/dml_provider_factory.cc | 6 +- .../cpu/tensor/quantize_linear_test.cc | 84 +++++++++++ .../providers/dml_safe_make_or_throw_test.cc | 139 ++++++++++++++++++ 22 files changed, 390 insertions(+), 104 deletions(-) create mode 100644 onnxruntime/core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h create mode 100644 onnxruntime/test/providers/dml_safe_make_or_throw_test.cc diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp index 353f698bb6f2c..076027dd3672f 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp @@ -504,7 +504,7 @@ HRESULT STDMETHODCALLTYPE AbiCustomRegistry::RegisterOperatorKernel( InferAndVerifyOutputSizes(node, &defaultAttributesCapture, shapeInferrerCapture.Get(), constantCpuInputCapture, constantInputGetter, inputShapesOverrides, *outputShapes); // Create the kernel while allowing input shape and output shape queries according to options - ComPtr kernelInfoWrapper = wil::MakeOrThrow( + ComPtr kernelInfoWrapper = Dml::SafeMakeOrThrow( &protoHelper, executionHandle, true, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/BucketizedBufferAllocator.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/BucketizedBufferAllocator.cpp index 18b4b4593f537..ed99ac0fc7fc2 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/BucketizedBufferAllocator.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/BucketizedBufferAllocator.cpp @@ -132,7 +132,7 @@ namespace Dml assert(resourceWrapper->GetD3D12Resource()->GetDesc().Width == bucketSize); assert(resourceWrapper != nullptr); - ComPtr allocInfo = wil::MakeOrThrow( + ComPtr allocInfo = Dml::SafeMakeOrThrow( this, ++m_currentAllocationId, resourceId, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommittedResourceAllocator.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommittedResourceAllocator.cpp index 54393e9bf1539..2934fd0c11516 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommittedResourceAllocator.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommittedResourceAllocator.cpp @@ -22,7 +22,7 @@ namespace Dml )); ComPtr resourceWrapper; - wil::MakeOrThrow(std::move(resource)).As(&resourceWrapper); + Dml::SafeMakeOrThrow(std::move(resource)).As(&resourceWrapper); return resourceWrapper; } } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlExternalBufferAllocator.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlExternalBufferAllocator.h index c99d686349e94..158c102d69ee7 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlExternalBufferAllocator.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlExternalBufferAllocator.h @@ -48,9 +48,9 @@ namespace Dml constexpr uint64_t pooledResourceId = 0; // Not a pooled resource Microsoft::WRL::ComPtr resourceWrapper; - wil::MakeOrThrow(std::move(resource)).As(&resourceWrapper); + Dml::SafeMakeOrThrow(std::move(resource)).As(&resourceWrapper); - Microsoft::WRL::ComPtr allocInfo = wil::MakeOrThrow( + Microsoft::WRL::ComPtr allocInfo = Dml::SafeMakeOrThrow( nullptr, 0, pooledResourceId, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index 6d8d5453b9fc0..cd7dfd46485af 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -55,7 +55,7 @@ namespace Dml _Out_ std::shared_ptr* registry, _Out_ std::shared_ptr* internalRegInfoMap) { - ComPtr abiRegistry = wil::MakeOrThrow(); + ComPtr abiRegistry = Dml::SafeMakeOrThrow(); Dml::RegisterDmlOperators(abiRegistry.Get()); assert(abiRegistry->GetRegistries().size() == 1); @@ -88,7 +88,7 @@ namespace Dml ComPtr device; GRAPHICS_THROW_IF_FAILED(dmlDevice->GetParentDevice(IID_GRAPHICS_PPV_ARGS(device.GetAddressOf()))); - m_impl = wil::MakeOrThrow(dmlDevice, device.Get(), executionContext, enableMetacommands, + m_impl = Dml::SafeMakeOrThrow(dmlDevice, device.Get(), executionContext, enableMetacommands, enableGraphCapture, enableSyncSpinning, disableMemoryArena); } @@ -1298,9 +1298,9 @@ namespace Dml uint64_t pooledResourceId = 0; // Not a pooled resource ComPtr resourceWrapper; - wil::MakeOrThrow(pResource).As(&resourceWrapper); + Dml::SafeMakeOrThrow(pResource).As(&resourceWrapper); - ComPtr allocInfo = wil::MakeOrThrow(nullptr, 0, pooledResourceId, resourceWrapper.Get(), (size_t)pResource->GetDesc().Width); + ComPtr allocInfo = Dml::SafeMakeOrThrow(nullptr, 0, pooledResourceId, resourceWrapper.Get(), (size_t)pResource->GetDesc().Width); return allocInfo.Detach(); } void FreeGPUAllocation(void* ptr) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp index 22de743f6e718..51c25d6d40c5b 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp @@ -291,7 +291,7 @@ namespace Dml::GraphDescBuilder if (iter != isInitializerTransferable.end()) { // Using const_cast here is simpler than making surrounding code const correct. - tensorWrapper = wil::MakeOrThrow(const_cast(iter->second.first), modelPath); + tensorWrapper = Dml::SafeMakeOrThrow(const_cast(iter->second.first), modelPath); } return tensorWrapper; }; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp index fe52f27b35bb8..13ce9afa99b1e 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp @@ -868,7 +868,7 @@ namespace Windows::AI::MachineLearning::Adapter const onnx::TensorProto* tensorProto = &attributeProto->t(); // An empty path is used as external weights are not currently supported in this case - Microsoft::WRL::ComPtr tensorWrapper = wil::MakeOrThrow(const_cast(tensorProto), std::filesystem::path()); + Microsoft::WRL::ComPtr tensorWrapper = Dml::SafeMakeOrThrow(const_cast(tensorProto), std::filesystem::path()); *tensor = tensorWrapper.Detach(); return S_OK; } @@ -1977,7 +1977,7 @@ namespace Windows::AI::MachineLearning::Adapter auto inputTensor = m_impl->Input(gsl::narrow_cast(inputIndex)); if (inputTensor != nullptr) { - ComPtr tensorWrapper = wil::MakeOrThrow( + ComPtr tensorWrapper = Dml::SafeMakeOrThrow( const_cast(inputTensor), IsAllocationInterface(inputTensor->Location()), m_winmlProvider.Get(), @@ -2019,7 +2019,7 @@ namespace Windows::AI::MachineLearning::Adapter auto elemTensor = const_cast(&inputTensorSeq->Get(sequenceIndex)); if (elemTensor != nullptr) { - ComPtr tensorWrapper = wil::MakeOrThrow( + ComPtr tensorWrapper = Dml::SafeMakeOrThrow( elemTensor, IsAllocationInterface(elemTensor->Location()), m_winmlProvider.Get(), @@ -2119,7 +2119,7 @@ namespace Windows::AI::MachineLearning::Adapter auto elemTensor = const_cast(&outputTensorSeq->Get(sequenceIndex)); if (elemTensor != nullptr) { - ComPtr tensorWrapper = wil::MakeOrThrow( + ComPtr tensorWrapper = Dml::SafeMakeOrThrow( elemTensor, IsAllocationInterface(elemTensor->Location()), m_winmlProvider.Get(), @@ -2212,7 +2212,7 @@ namespace Windows::AI::MachineLearning::Adapter auto outputTensor = m_impl->Output(outputIndex, shape); if (outputTensor) { - ComPtr tensorWrapper = wil::MakeOrThrow( + ComPtr tensorWrapper = Dml::SafeMakeOrThrow( const_cast(outputTensor), IsAllocationInterface(outputTensor->Location()), m_winmlProvider.Get(), @@ -2377,7 +2377,7 @@ namespace Windows::AI::MachineLearning::Adapter const onnxruntime::Tensor* tensor = nullptr; if (kerneInfo.TryGetConstantInput(index, &tensor)) { - tensorWrapper = wil::MakeOrThrow( + tensorWrapper = Dml::SafeMakeOrThrow( const_cast(tensor), IsAllocationInterface(tensor->Location()), winmlProviderCapture.Get(), @@ -2396,7 +2396,7 @@ namespace Windows::AI::MachineLearning::Adapter } // Create the kernel while allowing input shape and output shape queries according to options - ComPtr kernelInfoWrapper = wil::MakeOrThrow( + ComPtr kernelInfoWrapper = Dml::SafeMakeOrThrow( &kerneInfo, m_abiExecutionObject.Get(), nullptr, @@ -2443,7 +2443,7 @@ namespace Windows::AI::MachineLearning::Adapter const auto* tensor = context->Input(gsl::narrow_cast(index)); if (tensor != nullptr) { - tensorWrapper = wil::MakeOrThrow( + tensorWrapper = Dml::SafeMakeOrThrow( const_cast(tensor), IsAllocationInterface(tensor->Location()), winmlProviderCapture.Get(), @@ -2464,7 +2464,7 @@ namespace Windows::AI::MachineLearning::Adapter for (uint32_t sequenceIndex = 0; sequenceIndex < tensorSequence->Size(); ++sequenceIndex) { auto& tensor = tensorSequence->Get(sequenceIndex); - auto tensorWrapper = wil::MakeOrThrow( + auto tensorWrapper = Dml::SafeMakeOrThrow( const_cast(&tensor), IsAllocationInterface(tensor.Location()), winmlProviderCapture.Get(), @@ -2491,7 +2491,7 @@ namespace Windows::AI::MachineLearning::Adapter } // Create the kernel while allowing input shape and output shape queries according to options - ComPtr kernelInfoWrapper = wil::MakeOrThrow( + ComPtr kernelInfoWrapper = Dml::SafeMakeOrThrow( &Info(), m_abiExecutionObject.Get(), &inputShapes, @@ -2569,7 +2569,7 @@ namespace Windows::AI::MachineLearning::Adapter EdgeShapes localInferredOutputShapes; ComPtr localKernel = inferShapesAndCreateKernel(local_input_shapes, localInferredOutputShapes); - ComPtr kernelContextWrapper = wil::MakeOrThrow( + ComPtr kernelContextWrapper = Dml::SafeMakeOrThrow( context, Info().GetExecutionProvider(), m_internalOperator, @@ -2588,7 +2588,7 @@ namespace Windows::AI::MachineLearning::Adapter } } - ComPtr kernelContextWrapper = wil::MakeOrThrow( + ComPtr kernelContextWrapper = Dml::SafeMakeOrThrow( context, Info().GetExecutionProvider(), m_internalOperator, @@ -2811,7 +2811,7 @@ namespace Windows::AI::MachineLearning::Adapter onnxruntime::ProtoHelperNodeContext protoContext(node); onnxruntime::OpNodeProtoHelper info(&protoContext); - ComPtr inferenceContext = wil::MakeOrThrow(&info, inputShapes, outputShapes, defaultAttributes, requiredConstantCpuInputs, constantInputGetter); + ComPtr inferenceContext = Dml::SafeMakeOrThrow(&info, inputShapes, outputShapes, defaultAttributes, requiredConstantCpuInputs, constantInputGetter); outputShapes.Reset(info.GetOutputCount()); @@ -2865,13 +2865,13 @@ namespace Windows::AI::MachineLearning::Adapter [ctx](uint32_t index) { // An empty path is used as external weights are not currently supported in this case - Microsoft::WRL::ComPtr tensorWrapper = wil::MakeOrThrow( + Microsoft::WRL::ComPtr tensorWrapper = Dml::SafeMakeOrThrow( const_cast(ctx->getInputData(index)), std::filesystem::path()); return tensorWrapper; } ); - return wil::MakeOrThrow(info, ctx, requiredConstantCpuInputs, mlOperatorTensorGetter); + return Dml::SafeMakeOrThrow(info, ctx, requiredConstantCpuInputs, mlOperatorTensorGetter); } MLSchemaInferenceContext::MLSchemaInferenceContext( @@ -2952,7 +2952,7 @@ namespace Windows::AI::MachineLearning::Adapter const AttributeMap* defaultAttributes) { MLOperatorTensorGetter mLOperatorTensorGetter = MLOperatorTensorGetter(); - return wil::MakeOrThrow(info, defaultAttributes, mLOperatorTensorGetter); + return Dml::SafeMakeOrThrow(info, defaultAttributes, mLOperatorTensorGetter); } MLSupportQueryContext::MLSupportQueryContext( diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlDFT.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlDFT.h index 1de88a61a0d77..25210c146a6b6 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlDFT.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlDFT.h @@ -1097,7 +1097,7 @@ class GpuDFTOperatorFactory : public WRL::Base version = 20; } - auto dftOperator = wil::MakeOrThrow(context, version); + auto dftOperator = Dml::SafeMakeOrThrow(context, version); dftOperator.CopyTo(kernel); return S_OK; } @@ -1177,8 +1177,8 @@ class GpuDFTOperatorFactory : public WRL::Base kernelDescription.options = MLOperatorKernelOptions::None; kernelDescription.executionOptions = 0; - auto shareInferrer = wil::MakeOrThrow(); - auto factory = wil::MakeOrThrow(); + auto shareInferrer = Dml::SafeMakeOrThrow(); + auto factory = Dml::SafeMakeOrThrow(); std::array requiredConstantCpuInputs = { 1, 2 }; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlGridSample.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlGridSample.h index 5ba936ddf3976..6d7a089103c9b 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlGridSample.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlGridSample.h @@ -747,7 +747,7 @@ class DmlGridSampleOperatorFactory : public WRL::Base { try { - auto dftOperator = wil::MakeOrThrow(context); + auto dftOperator = Dml::SafeMakeOrThrow(context); dftOperator.CopyTo(kernel); return S_OK; } @@ -832,8 +832,8 @@ class DmlGridSampleOperatorFactory : public WRL::Base kernelDescription.options = MLOperatorKernelOptions::None; kernelDescription.executionOptions = 0; - auto shareInferrer = wil::MakeOrThrow(); - auto factory = wil::MakeOrThrow(); + auto shareInferrer = Dml::SafeMakeOrThrow(); + auto factory = Dml::SafeMakeOrThrow(); ComPtr registryPrivate; ORT_THROW_IF_FAILED(registry->QueryInterface(IID_PPV_ARGS(®istryPrivate))); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp index 287f1e5b6dfe7..2ee85b01a9a2e 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp @@ -907,4 +907,71 @@ namespace Dml bufferTensorDesc->TotalTensorSizeInBytes = (elementSize + 3) & ~3; } + void DmlOperator::BroadcastQuantizationParameters( + const MLOperatorKernelCreationContext& kernelInfo, + gsl::span outputShape + ) + { + const uint32_t outputShapeDimCount = gsl::narrow_cast(outputShape.size()); + + uint32_t axis = 0; + + // If an axis was explicitly passed (or the default value 1 is set from the schema), + // then other inputs are broadcasting to the shape of the input data tensor. + if (kernelInfo.HasAttribute(AttrName::Axis, MLOperatorAttributeType::Int)) + { + // Avoid validating the axis until later because the axis parameter is ignorable unless + // broadcasting is actually needed. ONNX opset 13 returns a default value of 1 for the + // "axis" attribute even when the attribute doesn't actually exist in the model, which + // would cause a validation failure here. + const int32_t signedAxis = gsl::narrow_cast(kernelInfo.GetAttribute(AttrName::Axis)); + axis = Dml::HandleNegativeAxis(signedAxis, outputShapeDimCount, /*validateAxis*/ false); + } + + // Explicitly reshape each of the inputs after the first input (scale tensor and optional zero point tensor). + for (uint32_t index = 1, inputCount = gsl::narrow_cast(m_inputTensorDescs.size()); index < inputCount; ++index) + { + if (!kernelInfo.IsInputValid(index)) + { + continue; + } + + auto edgeDesc = kernelInfo.GetInputEdgeDescription(index); + assert(edgeDesc.edgeType == MLOperatorEdgeType::Tensor); + + // Fix up the tensor shape by filling with trailing ones. So input[2,3] with axis=0 and scale[2] + // becomes scale[2,1], so that broadcasting works correctly. + std::vector inputTensorShape = kernelInfo.GetTensorShapeDescription().GetInputTensorShape(index); + + // If the input tensor is a 1D vector, then extra massaging is needed to project their + // 1D vectors back to the full shape for broadcasting along the given axis. + // The 1D vector should have a length equal to the output tensor's dimension on that axis. + if (inputTensorShape.size() == 1 && inputTensorShape != std::vector(outputShape.begin(), outputShape.end())) + { + ML_CHECK_VALID_ARGUMENT(axis < outputShapeDimCount); + uint32_t broadcastAxisLength = outputShape[axis]; + ML_CHECK_VALID_ARGUMENT( + (inputTensorShape[0] == broadcastAxisLength) || + // Treat as broadcast dimension to match CPU behavior. + (inputTensorShape[0] == 1) + ); + inputTensorShape.insert(inputTensorShape.begin(), axis, 1); + inputTensorShape.insert(inputTensorShape.end(), outputShapeDimCount - 1 - axis, 1); + } + // For any other shape (scalar/ND), leave it alone, and the TensorDesc constructor + // will apply broadcasting with standard elementwise alignment. + + m_inputTensorDescs[index] = TensorDesc( + edgeDesc.tensorDataType, + outputShape, + gsl::make_span(inputTensorShape), + TensorAxis::DoNotCoerce, + TensorAxis::W, + TensorAxis::RightAligned, + NchwDimensionCount, // minDimensionCount + 0 // guaranteedBaseOffsetAlignment + ); + } + } + } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h index fa54d4b041b5f..002541e23c47c 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h @@ -149,6 +149,15 @@ namespace Dml uint32_t minDimensionCount = NchwDimensionCount ) const; + // Reshapes scale and zero_point tensor descriptors (inputs after index 0) so that their + // dimension count matches the output shape, enabling correct broadcasting in DML. + // For 1D per-axis tensors, the shape is projected along the given axis (e.g. scale[6] + // with axis=0 on a 5D output becomes [6,1,1,1,1]). + void BroadcastQuantizationParameters( + const MLOperatorKernelCreationContext& kernelInfo, + gsl::span outputShape + ); + static void TryConvertTensorToBroadcastScalar( const MLOperatorKernelCreationContext& kernelInfo, const DML_TENSOR_DESC* tensor, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp index d4d7ee1311874..b64a5265f56e3 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp @@ -542,64 +542,7 @@ class DmlOperatorElementwiseQLinear : public DmlOperator const DML_TENSOR_DATA_TYPE outputDataType = m_outputTensorDescs[0].GetDmlDataType(); bool hasZeroPointTensor = kernelInfo.IsInputValid(2); - uint32_t axis = 0; - - // If an axis was given explicitly passed (or the default value 1 is set from the schema), - // then other inputs are broadcasting to the shape of the input data tensor. - if (kernelInfo.HasAttribute(AttrName::Axis, MLOperatorAttributeType::Int)) - { - // Avoid validating the axis until later because the axis parameter is ignorable unless - // broadcasting is actually needed. ONNX opset 13 returns a default value of 1 for the - // "axis" attribute even when the attribute doesn't actually exist in the model, which - // would cause a validation failure here. - const int32_t signedAxis = gsl::narrow_cast(kernelInfo.GetAttribute(AttrName::Axis)); - axis = Dml::HandleNegativeAxis(signedAxis, outputShapeDimCount, /*validateAxis*/ false); - } - - // Explicitly reshape each of the inputs after the first input (scale tensor and optional zero point tensor). - for (uint32_t index = 1, inputCount = gsl::narrow_cast(m_inputTensorDescs.size()); index < inputCount; ++index) - { - if (!kernelInfo.IsInputValid(index)) - { - continue; - } - - auto edgeDesc = kernelInfo.GetInputEdgeDescription(index); - assert(edgeDesc.edgeType == MLOperatorEdgeType::Tensor); - - // Fix up the the tensor shape by filling with trailing ones. So input[2,3] with axis=0 and scale[2] - // becomes scale[2,1], so that broadcasting works correctly. - std::vector inputTensorShape = kernelInfo.GetTensorShapeDescription().GetInputTensorShape(index); - - // If the input tensor is a 1D vector, then extra massaging is needed to project their - // 1D vectors back to the full shape for broadcasting along the given axis. - // The 1D vector should have a length equal to the output tensor's dimension on that axis. - if (inputTensorShape.size() == 1 && inputTensorShape != outputShape) - { - ML_CHECK_VALID_ARGUMENT(axis < outputShapeDimCount); - uint32_t broadcastAxisLength = outputShape[axis]; - ML_CHECK_VALID_ARGUMENT( - (inputTensorShape[0] == broadcastAxisLength) || - // Treat as broadcast dimension to match CPU behavior. - (inputTensorShape[0] == 1) - ); - inputTensorShape.insert(inputTensorShape.begin(), axis, 1); - inputTensorShape.insert(inputTensorShape.end(), outputShapeDimCount - 1 - axis, 1); - } - // For any other shape (scalar/ND), leave it alone, and the TensorDesc constructor - // will apply broadcasting with standard elementwise alignment. - - m_inputTensorDescs[index] = TensorDesc( - edgeDesc.tensorDataType, - gsl::make_span(outputShape), - gsl::make_span(inputTensorShape), - TensorAxis::DoNotCoerce, - TensorAxis::W, - TensorAxis::RightAligned, - NchwDimensionCount, // minDimensionCount - 0 // guaranteedBaseOffsetAlignment - ); - } + BroadcastQuantizationParameters(kernelInfo, gsl::make_span(outputShape)); std::vector inputDescs = GetDmlInputDescs(); std::vector outputDescs = GetDmlOutputDescs(); @@ -630,6 +573,8 @@ class DmlOperatorQuantization21 : public DmlOperator const DML_TENSOR_DATA_TYPE outputDataType = m_outputTensorDescs[0].GetDmlDataType(); bool hasZeroPointTensor = kernelInfo.IsInputValid(2); + BroadcastQuantizationParameters(kernelInfo, gsl::make_span(outputShape)); + std::vector inputDescs = GetDmlInputDescs(); std::vector outputDescs = GetDmlOutputDescs(); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorNonZero.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorNonZero.cpp index bc29256dd2e28..83e35ae89282d 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorNonZero.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorNonZero.cpp @@ -76,7 +76,7 @@ class DmlOperatorNonZero: public DmlOperator // Create the DML output tensor for the number of nonzero elements onnxruntime::Tensor outputCountDml(onnxruntime::DataTypeImpl::GetType(), m_outputCountShape, executionProvider->GetGpuAllocator()); - Microsoft::WRL::ComPtr outputCountDmlWrapper = wil::MakeOrThrow( + Microsoft::WRL::ComPtr outputCountDmlWrapper = Dml::SafeMakeOrThrow( &outputCountDml, true, executionProvider, @@ -84,7 +84,7 @@ class DmlOperatorNonZero: public DmlOperator // Create the DML output tensor for the coordinates (not cropped) onnxruntime::Tensor intermediateCoordinatesDml(onnxruntime::DataTypeImpl::GetType(), m_outputCoordinatesShape, executionProvider->GetGpuAllocator()); - Microsoft::WRL::ComPtr intermediateCoordinatesDmlWrapper = wil::MakeOrThrow( + Microsoft::WRL::ComPtr intermediateCoordinatesDmlWrapper = Dml::SafeMakeOrThrow( &intermediateCoordinatesDml, true, executionProvider, @@ -105,7 +105,7 @@ class DmlOperatorNonZero: public DmlOperator // Copy the number of nonzero elements back to the CPU onnxruntime::Tensor outputCountCpu(onnxruntime::DataTypeImpl::GetType(), {1}, executionProvider->GetCpuInputAllocator()); - Microsoft::WRL::ComPtr outputCountCpuWrapper = wil::MakeOrThrow( + Microsoft::WRL::ComPtr outputCountCpuWrapper = Dml::SafeMakeOrThrow( &outputCountCpu, false, executionProvider, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlSTFT.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlSTFT.h index e2f38231f7295..091a82daefbdc 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlSTFT.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlSTFT.h @@ -238,7 +238,7 @@ class DmlSTFTOperator : public WRL::Base constexpr uint32_t dftAxis = 1; constexpr bool dftIsInverse = false; - m_dftOperator.op = wil::MakeOrThrow( + m_dftOperator.op = Dml::SafeMakeOrThrow( m_d3dDevice.Get(), dftAxis, params.isOnesided, @@ -516,7 +516,7 @@ class DmlSTFTOperatorFactory : public WRL::Base { try { - auto dftOperator = wil::MakeOrThrow(context); + auto dftOperator = Dml::SafeMakeOrThrow(context); dftOperator.CopyTo(kernel); return S_OK; } @@ -574,8 +574,8 @@ class DmlSTFTOperatorFactory : public WRL::Base kernelDescription.options = MLOperatorKernelOptions::None; kernelDescription.executionOptions = 0; - auto shareInferrer = wil::MakeOrThrow(); - auto factory = wil::MakeOrThrow(); + auto shareInferrer = Dml::SafeMakeOrThrow(); + auto factory = Dml::SafeMakeOrThrow(); std::array requiredConstantCpuInputs = { /*frame_step*/1, /*frame_length*/3 }; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp index b0b37d01370bc..26f998c7521a2 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -1314,18 +1314,18 @@ void RegisterDmlOperators(IMLOperatorRegistry* registry) totalTypeCount += typeConstraints[i].allowedTypeCount; } - ComPtr factory = wil::MakeOrThrow(information.creationFunction); + ComPtr factory = Dml::SafeMakeOrThrow(information.creationFunction); ComPtr shapeInferrer; if (information.shapeInferenceFunction) { - shapeInferrer = wil::MakeOrThrow(information.shapeInferenceFunction); + shapeInferrer = Dml::SafeMakeOrThrow(information.shapeInferenceFunction); } ComPtr supportQuery; if (information.supportQueryFunction) { - supportQuery = wil::MakeOrThrow(information.supportQueryFunction); + supportQuery = Dml::SafeMakeOrThrow(information.supportQueryFunction); } ORT_THROW_IF_FAILED(registryPrivate->RegisterOperatorKernel( diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h new file mode 100644 index 0000000000000..c2740470cbc0a --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +// Drop-in replacement for wil::MakeOrThrow that avoids an ASan false positive. +// WRL's MakeAllocator stores its buffer as char*, so if the constructor throws, +// ~MakeAllocator calls delete on a char* — passing sizeof(char)=1 to sized +// operator delete instead of sizeof(T). With the default MSVC allocator, this is +// benign (sized delete ignores the size), but ASan flags it as +// new-delete-type-mismatch. This helper uses placement new with correctly-sized +// cleanup to avoid the issue. +namespace Dml +{ + template + Microsoft::WRL::ComPtr SafeMakeOrThrow(TArgs&&... args) + { + void* buffer = ::operator new(sizeof(T)); + T* raw = nullptr; + try + { + raw = new (buffer) T(std::forward(args)...); + } + catch (...) + { + ::operator delete(buffer, sizeof(T)); + throw; + } + Microsoft::WRL::ComPtr result; + result.Attach(raw); + return result; + } +} // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/precomp.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/precomp.h index e9df3fd20aff9..b9febb8171e0d 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/precomp.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/precomp.h @@ -25,6 +25,7 @@ #include #include +#include "SafeMakeOrThrow.h" #include diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h index ac77616cb96f0..dec84d9945569 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h @@ -5,6 +5,7 @@ #include "core/providers/dml/DmlExecutionProvider/inc/MLOperatorAuthor.h" #include "MLOperatorAuthorPrivate.h" +#include "core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h" #include "core/framework/int4.h" #include #include @@ -972,7 +973,7 @@ class MLOperatorKernel : public Microsoft::WRL::RuntimeClass< { ORT_TRY { - Microsoft::WRL::ComPtr kernel = wil::MakeOrThrow(MLOperatorKernelCreationContext(&info)); + Microsoft::WRL::ComPtr kernel = Dml::SafeMakeOrThrow(MLOperatorKernelCreationContext(&info)); *opKernel = kernel.Detach(); return S_OK; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/SchemaInferenceOverrider.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/SchemaInferenceOverrider.h index fa04bcf6edf41..597780a9f448b 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/SchemaInferenceOverrider.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/SchemaInferenceOverrider.h @@ -5,6 +5,7 @@ #include "OperatorHelper.h" #include "OperatorVersions.h" +#include "core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h" namespace SchemaInferenceOverrider { @@ -21,7 +22,7 @@ namespace SchemaInferenceOverrider ) { Microsoft::WRL::ComPtr shapeInferrer = - wil::MakeOrThrow(OperatorHelper::ShapeInferenceFunction); + Dml::SafeMakeOrThrow(OperatorHelper::ShapeInferenceFunction); auto schema = const_cast(onnx::OpSchemaRegistry::Schema(name, version)); diff --git a/onnxruntime/core/providers/dml/dml_provider_factory.cc b/onnxruntime/core/providers/dml/dml_provider_factory.cc index c72ce205e5fbb..c0ddc44d0ca57 100644 --- a/onnxruntime/core/providers/dml/dml_provider_factory.cc +++ b/onnxruntime/core/providers/dml/dml_provider_factory.cc @@ -21,6 +21,8 @@ using Microsoft::WRL::ComPtr; #include #include +#include "core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h" + #include "core/providers/dml/dml_provider_factory.h" #include "core/providers/dml/dml_provider_factory_creator.h" #include "core/session/abi_session_options_impl.h" @@ -89,11 +91,11 @@ std::unique_ptr DMLProviderFactory::CreateProvider() { // First, check if an I/O binding API that was used before this session or another session has already created a queue if (FAILED(d3d12_device->GetPrivateData(dml_execution_context_guid, &execution_context_ptr_size, execution_context.GetAddressOf()))) { - execution_context = wil::MakeOrThrow(d3d12_device.Get(), dml_device_.Get(), cmd_queue_.Get(), true, true); + execution_context = Dml::SafeMakeOrThrow(d3d12_device.Get(), dml_device_.Get(), cmd_queue_.Get(), true, true); ORT_THROW_IF_FAILED(d3d12_device->SetPrivateDataInterface(dml_execution_context_guid, execution_context.Get())); } } else { - execution_context = wil::MakeOrThrow(d3d12_device.Get(), dml_device_.Get(), cmd_queue_.Get(), cpu_sync_spinning_enabled_, false); + execution_context = Dml::SafeMakeOrThrow(d3d12_device.Get(), dml_device_.Get(), cmd_queue_.Get(), cpu_sync_spinning_enabled_, false); } auto provider = Dml::CreateExecutionProvider(dml_device_.Get(), execution_context.Get(), metacommands_enabled_, graph_capture_enabled_, cpu_sync_spinning_enabled_, disable_memory_arena_); diff --git a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc index 62fe0ffbc0581..6787efc11517e 100644 --- a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc @@ -531,6 +531,90 @@ TEST(QuantizeLinearOpTest, Int8) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); } +// Repro for new-delete-type-mismatch in DML EP during graph fusion. +// QuantizeLinear float32→int8 with 5D input triggers a type-size +// mismatch (192 bytes allocated, 1 byte deallocated) visible under ASan. +TEST(QuantizeLinearOpTest, Int8_5D_DML_TypeMismatch) { + auto dml_ep = DefaultDmlExecutionProvider(); + if (!dml_ep) { + GTEST_SKIP() << "Skipping because DML EP is not available."; + } + + OpTester test("QuantizeLinear", 13); + std::vector dims{6, 1, 1, 1, 1}; + test.AddInput("x", dims, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); + test.AddInput("y_scale", {}, {1.0f}); + test.AddInput("y_zero_point", {}, {0}); + test.AddOutput("y", dims, {1, 2, 3, 4, 5, 6}); + + std::vector> execution_providers; + execution_providers.emplace_back(std::move(dml_ep)); + test.ConfigEps(std::move(execution_providers)) + .RunWithConfig(); +} + +// Same as above but with per-axis quantization along axis 0 to exercise +// the DML graph fusion path with per-channel int8 quantization. +TEST(QuantizeLinearOpTest, Int8_5D_PerAxis_DML_TypeMismatch) { + auto dml_ep = DefaultDmlExecutionProvider(); + if (!dml_ep) { + GTEST_SKIP() << "Skipping because DML EP is not available."; + } + + OpTester test("QuantizeLinear", 13); + std::vector dims{6, 1, 1, 1, 1}; + test.AddAttribute("axis", 0); + test.AddInput("x", dims, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); + test.AddInput("y_scale", {6}, {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f}); + test.AddInput("y_zero_point", {6}, {0, 0, 0, 0, 0, 0}); + test.AddOutput("y", dims, {1, 2, 3, 4, 5, 6}); + + std::vector> execution_providers; + execution_providers.emplace_back(std::move(dml_ep)); + test.ConfigEps(std::move(execution_providers)) + .RunWithConfig(); +} + +// Opset 21 QuantizeLinear float32→uint8 WITHOUT zero_point. +// Without zero_point, the output type defaults to uint8. +TEST(QuantizeLinearOpTest, Uint8_5D_NoZeroPoint_Opset21_DML) { + auto dml_ep = DefaultDmlExecutionProvider(); + if (!dml_ep) { + GTEST_SKIP() << "Skipping because DML EP is not available."; + } + + OpTester test("QuantizeLinear", 21); + std::vector dims{6, 1, 1, 1, 1}; + test.AddInput("x", dims, {0.0f, 51.0f, 102.0f, 153.0f, 204.0f, 255.0f}); + test.AddInput("y_scale", {}, {1.0f}); + test.AddOutput("y", dims, {0, 51, 102, 153, 204, 255}); + + std::vector> execution_providers; + execution_providers.emplace_back(std::move(dml_ep)); + test.ConfigEps(std::move(execution_providers)) + .RunWithConfig(); +} + +// Opset 21 QuantizeLinear float32→int8 with zero_point (the customer's exact scenario). +TEST(QuantizeLinearOpTest, Int8_5D_WithZeroPoint_Opset21_DML) { + auto dml_ep = DefaultDmlExecutionProvider(); + if (!dml_ep) { + GTEST_SKIP() << "Skipping because DML EP is not available."; + } + + OpTester test("QuantizeLinear", 21); + std::vector dims{6, 1, 1, 1, 1}; + test.AddInput("x", dims, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); + test.AddInput("y_scale", {}, {1.0f}); + test.AddInput("y_zero_point", {}, {0}); + test.AddOutput("y", dims, {1, 2, 3, 4, 5, 6}); + + std::vector> execution_providers; + execution_providers.emplace_back(std::move(dml_ep)); + test.ConfigEps(std::move(execution_providers)) + .RunWithConfig(); +} + // Test uint16 QuantizeLinear (per tensor) TEST(QuantizeLinearOpTest, Uint16) { OpTester test("QuantizeLinear", 21); diff --git a/onnxruntime/test/providers/dml_safe_make_or_throw_test.cc b/onnxruntime/test/providers/dml_safe_make_or_throw_test.cc new file mode 100644 index 0000000000000..8041f0dae8c28 --- /dev/null +++ b/onnxruntime/test/providers/dml_safe_make_or_throw_test.cc @@ -0,0 +1,139 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#ifdef USE_DML + +#include "gtest/gtest.h" + +#include +#include +#include "core/providers/dml/DmlExecutionProvider/src/SafeMakeOrThrow.h" + +#include + +namespace onnxruntime { +namespace test { + +// A trivial COM interface for testing. +MIDL_INTERFACE("A1B2C3D4-E5F6-7890-ABCD-EF1234567890") +ITestInterface : public IUnknown { + virtual int STDMETHODCALLTYPE GetValue() = 0; +}; + +// A RuntimeClass whose constructor succeeds and stores a value. +class SucceedingClass : public Microsoft::WRL::RuntimeClass< + Microsoft::WRL::RuntimeClassFlags, ITestInterface> { + public: + int value; + + SucceedingClass(int v) : value(v) {} + + int STDMETHODCALLTYPE GetValue() override { return value; } +}; + +// A RuntimeClass that tracks whether its destructor ran. +class TrackedClass : public Microsoft::WRL::RuntimeClass< + Microsoft::WRL::RuntimeClassFlags, ITestInterface> { + public: + bool& destroyed; + + TrackedClass(bool& flag) : destroyed(flag) { destroyed = false; } + ~TrackedClass() { destroyed = true; } + + int STDMETHODCALLTYPE GetValue() override { return 42; } +}; + +// A RuntimeClass whose constructor always throws. +// Uses a ref-counted witness to verify cleanup: the witness is destroyed +// (via Release) during stack unwinding if memory is freed correctly. +class ThrowingClass : public Microsoft::WRL::RuntimeClass< + Microsoft::WRL::RuntimeClassFlags, ITestInterface> { + public: + Microsoft::WRL::ComPtr witness; + + ThrowingClass(bool& witness_destroyed) { + // Create a witness that will be destroyed when this object's members + // are cleaned up during stack unwinding. + witness = Dml::SafeMakeOrThrow(witness_destroyed); + throw std::runtime_error("intentional throw"); + } + + int STDMETHODCALLTYPE GetValue() override { return -1; } +}; + +// Verify that SafeMakeOrThrow creates an object with ref count 1, +// and that the object is properly released when the ComPtr goes out of scope. +TEST(SafeMakeOrThrowTest, SuccessPath_RefCountIsOne) { + Microsoft::WRL::ComPtr obj = Dml::SafeMakeOrThrow(123); + + ASSERT_NE(obj.Get(), nullptr); + EXPECT_EQ(obj->GetValue(), 123); + + // AddRef/Release to observe ref count: AddRef returns new count. + unsigned long refAfterAdd = obj->AddRef(); + EXPECT_EQ(refAfterAdd, 2u); + + unsigned long refAfterRelease = obj->Release(); + EXPECT_EQ(refAfterRelease, 1u); +} + +// Verify that the object is destroyed when the last ComPtr releases it. +TEST(SafeMakeOrThrowTest, SuccessPath_DestructorRunsOnRelease) { + bool destroyed = false; + { + auto obj = Dml::SafeMakeOrThrow(destroyed); + EXPECT_FALSE(destroyed); + } + // ComPtr went out of scope — destructor should have run. + EXPECT_TRUE(destroyed); +} + +// Verify that copying the ComPtr increments the ref count and +// the object survives until the last reference is released. +TEST(SafeMakeOrThrowTest, SuccessPath_MultipleReferences) { + bool destroyed = false; + Microsoft::WRL::ComPtr copy; + { + auto obj = Dml::SafeMakeOrThrow(destroyed); + copy = obj; + EXPECT_FALSE(destroyed); + } + // Original ComPtr gone, but copy still holds a reference. + EXPECT_FALSE(destroyed); + + copy.Reset(); + EXPECT_TRUE(destroyed); +} + +// Verify that when the constructor throws, the exception propagates +// and sub-objects are properly cleaned up (no leak). +TEST(SafeMakeOrThrowTest, FailurePath_ConstructorThrows) { + bool witness_destroyed = false; + EXPECT_THROW( + Dml::SafeMakeOrThrow(witness_destroyed), + std::runtime_error); + // The witness ComPtr member was constructed before the throw. + // If cleanup worked correctly, the witness should have been destroyed + // when the ThrowingClass sub-objects were unwound. + EXPECT_TRUE(witness_destroyed); +} + +// Verify that QI works correctly on a SafeMakeOrThrow-created object. +TEST(SafeMakeOrThrowTest, SuccessPath_QueryInterface) { + auto obj = Dml::SafeMakeOrThrow(42); + + Microsoft::WRL::ComPtr unk; + HRESULT hr = obj.As(&unk); + EXPECT_EQ(hr, S_OK); + EXPECT_NE(unk.Get(), nullptr); + + Microsoft::WRL::ComPtr iface; + hr = unk.As(&iface); + EXPECT_EQ(hr, S_OK); + EXPECT_EQ(iface->GetValue(), 42); +} + +} // namespace test +} // namespace onnxruntime + +#endif // USE_DML From 0730bc5581c5eadca399e57ca28b0a037d4ac8a8 Mon Sep 17 00:00:00 2001 From: adrastogi Date: Thu, 16 Apr 2026 16:55:45 -0700 Subject: [PATCH 4/4] Fix AddInitializers buffer size for sub-byte types --- onnxruntime/test/unittest_util/base_tester.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/unittest_util/base_tester.cc b/onnxruntime/test/unittest_util/base_tester.cc index 2e0459103a7c9..4d1b2a210599a 100644 --- a/onnxruntime/test/unittest_util/base_tester.cc +++ b/onnxruntime/test/unittest_util/base_tester.cc @@ -74,7 +74,9 @@ void BaseTester::AddInitializers(onnxruntime::Graph& graph) { tensor_proto.add_string_data(string_data[i]); } } else { - auto buffer_size = tensor.DataType()->Size() * shape.Size(); + // Note: need to use Tensor::CalculateTensorStorageSize (instead of shape.Size() * elem_size) to properly + // calculate the storage size for sub-byte types (e.g., Int4 or Int2) + auto buffer_size = Tensor::CalculateTensorStorageSize(tensor.DataType(), shape); utils::SetRawDataInTensorProto(tensor_proto, tensor.DataRaw(), buffer_size); }