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/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/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/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/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/providers/cpu/tensor/quantize_linear_test.cc b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc index bf632d0b3bc40..6787efc11517e 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); @@ -490,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 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); + } +} 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); }