Fix new-delete mismatch in DML EP's QuantizeLinear operator

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<DmlGraphOpKernelInfoWrapper> kernelInfoWrapper = wil::MakeOrThrow<DmlGraphOpKernelInfoWrapper>(
+                    ComPtr<DmlGraphOpKernelInfoWrapper> kernelInfoWrapper = Dml::SafeMakeOrThrow<DmlGraphOpKernelInfoWrapper>(
                             &protoHelper,
                             executionHandle,
                             true,

onnxruntime/core/providers/dml/DmlExecutionProvider/src/BucketizedBufferAllocator.cpp

@@ -132,7 +132,7 @@ namespace Dml
         assert(resourceWrapper->GetD3D12Resource()->GetDesc().Width == bucketSize);
         assert(resourceWrapper != nullptr);
 
-        ComPtr<AllocationInfo> allocInfo = wil::MakeOrThrow<AllocationInfo>(
+        ComPtr<AllocationInfo> allocInfo = Dml::SafeMakeOrThrow<AllocationInfo>(
             this,
             ++m_currentAllocationId,
             resourceId,

onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommittedResourceAllocator.cpp

@@ -22,7 +22,7 @@ namespace Dml
         ));
 
         ComPtr<DmlResourceWrapper> resourceWrapper;
-        wil::MakeOrThrow<DmlCommittedResourceWrapper>(std::move(resource)).As(&resourceWrapper);
+        Dml::SafeMakeOrThrow<DmlCommittedResourceWrapper>(std::move(resource)).As(&resourceWrapper);
         return resourceWrapper;
     }
 }

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<DmlResourceWrapper> resourceWrapper;
-            wil::MakeOrThrow<DmlCommittedResourceWrapper>(std::move(resource)).As(&resourceWrapper);
+            Dml::SafeMakeOrThrow<DmlCommittedResourceWrapper>(std::move(resource)).As(&resourceWrapper);
 
-            Microsoft::WRL::ComPtr<AllocationInfo> allocInfo = wil::MakeOrThrow<AllocationInfo>(
+            Microsoft::WRL::ComPtr<AllocationInfo> allocInfo = Dml::SafeMakeOrThrow<AllocationInfo>(
                 nullptr,
                 0,
                 pooledResourceId,

onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp

@@ -55,7 +55,7 @@ namespace Dml
         _Out_ std::shared_ptr<onnxruntime::KernelRegistry>* registry,
         _Out_ std::shared_ptr<const InternalRegistrationInfoMap>* internalRegInfoMap)
     {
-        ComPtr<AbiCustomRegistry> abiRegistry = wil::MakeOrThrow<AbiCustomRegistry>();
+        ComPtr<AbiCustomRegistry> abiRegistry = Dml::SafeMakeOrThrow<AbiCustomRegistry>();
         Dml::RegisterDmlOperators(abiRegistry.Get());
 
         assert(abiRegistry->GetRegistries().size() == 1);
@@ -88,7 +88,7 @@ namespace Dml
         ComPtr<ID3D12Device> device;
         GRAPHICS_THROW_IF_FAILED(dmlDevice->GetParentDevice(IID_GRAPHICS_PPV_ARGS(device.GetAddressOf())));
 
-        m_impl = wil::MakeOrThrow<ExecutionProviderImpl>(dmlDevice, device.Get(), executionContext, enableMetacommands,
+        m_impl = Dml::SafeMakeOrThrow<ExecutionProviderImpl>(dmlDevice, device.Get(), executionContext, enableMetacommands,
                                                          enableGraphCapture, enableSyncSpinning, disableMemoryArena);
     }
 
@@ -1298,9 +1298,9 @@ namespace Dml
         uint64_t pooledResourceId = 0; // Not a pooled resource
 
         ComPtr<DmlResourceWrapper> resourceWrapper;
-        wil::MakeOrThrow<DmlCommittedResourceWrapper>(pResource).As(&resourceWrapper);
+        Dml::SafeMakeOrThrow<DmlCommittedResourceWrapper>(pResource).As(&resourceWrapper);
 
-        ComPtr<AllocationInfo> allocInfo = wil::MakeOrThrow<AllocationInfo>(nullptr, 0, pooledResourceId, resourceWrapper.Get(), (size_t)pResource->GetDesc().Width);
+        ComPtr<AllocationInfo> allocInfo = Dml::SafeMakeOrThrow<AllocationInfo>(nullptr, 0, pooledResourceId, resourceWrapper.Get(), (size_t)pResource->GetDesc().Width);
         return allocInfo.Detach();
     }
     void FreeGPUAllocation(void* ptr)

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<OnnxTensorWrapper>(const_cast<ONNX_NAMESPACE::TensorProto*>(iter->second.first), modelPath);
+                tensorWrapper = Dml::SafeMakeOrThrow<OnnxTensorWrapper>(const_cast<ONNX_NAMESPACE::TensorProto*>(iter->second.first), modelPath);
             }
             return tensorWrapper;
         };

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<IMLOperatorTensor> tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(const_cast<onnx::TensorProto*>(tensorProto), std::filesystem::path());
+              Microsoft::WRL::ComPtr<IMLOperatorTensor> tensorWrapper = Dml::SafeMakeOrThrow<OnnxTensorWrapper>(const_cast<onnx::TensorProto*>(tensorProto), std::filesystem::path());
               *tensor = tensorWrapper.Detach();
               return S_OK;
             }
@@ -1977,7 +1977,7 @@ namespace Windows::AI::MachineLearning::Adapter
                 auto inputTensor = m_impl->Input<onnxruntime::Tensor>(gsl::narrow_cast<int>(inputIndex));
                 if (inputTensor != nullptr)
                 {
-                    ComPtr<TensorWrapper> tensorWrapper = wil::MakeOrThrow<TensorWrapper>(
+                    ComPtr<TensorWrapper> tensorWrapper = Dml::SafeMakeOrThrow<TensorWrapper>(
                         const_cast<onnxruntime::Tensor*>(inputTensor),
                         IsAllocationInterface(inputTensor->Location()),
                         m_winmlProvider.Get(),
@@ -2019,7 +2019,7 @@ namespace Windows::AI::MachineLearning::Adapter
                 auto elemTensor = const_cast<onnxruntime::Tensor*>(&inputTensorSeq->Get(sequenceIndex));
                 if (elemTensor != nullptr)
                 {
-                    ComPtr<TensorWrapper> tensorWrapper = wil::MakeOrThrow<TensorWrapper>(
+                    ComPtr<TensorWrapper> tensorWrapper = Dml::SafeMakeOrThrow<TensorWrapper>(
                         elemTensor,
                         IsAllocationInterface(elemTensor->Location()),
                         m_winmlProvider.Get(),
@@ -2119,7 +2119,7 @@ namespace Windows::AI::MachineLearning::Adapter
                 auto elemTensor = const_cast<onnxruntime::Tensor*>(&outputTensorSeq->Get(sequenceIndex));
                 if (elemTensor != nullptr)
                 {
-                    ComPtr<TensorWrapper> tensorWrapper = wil::MakeOrThrow<TensorWrapper>(
+                    ComPtr<TensorWrapper> tensorWrapper = Dml::SafeMakeOrThrow<TensorWrapper>(
                         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> tensorWrapper = wil::MakeOrThrow<TensorWrapper>(
+                    ComPtr<TensorWrapper> tensorWrapper = Dml::SafeMakeOrThrow<TensorWrapper>(
                         const_cast<onnxruntime::Tensor*>(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>(
+                    tensorWrapper = Dml::SafeMakeOrThrow<TensorWrapper>(
                         const_cast<onnxruntime::Tensor*>(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<OpKernelInfoWrapper> kernelInfoWrapper = wil::MakeOrThrow<OpKernelInfoWrapper>(
+            ComPtr<OpKernelInfoWrapper> kernelInfoWrapper = Dml::SafeMakeOrThrow<OpKernelInfoWrapper>(
                 &kerneInfo,
                 m_abiExecutionObject.Get(),
                 nullptr,
@@ -2443,7 +2443,7 @@ namespace Windows::AI::MachineLearning::Adapter
                     const auto* tensor = context->Input<onnxruntime::Tensor>(gsl::narrow_cast<int>(index));
                     if (tensor != nullptr)
                     {
-                        tensorWrapper = wil::MakeOrThrow<TensorWrapper>(
+                        tensorWrapper = Dml::SafeMakeOrThrow<TensorWrapper>(
                             const_cast<onnxruntime::Tensor*>(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<TensorWrapper>(
+                            auto tensorWrapper = Dml::SafeMakeOrThrow<TensorWrapper>(
                                 const_cast<onnxruntime::Tensor*>(&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<OpKernelInfoWrapper> kernelInfoWrapper = wil::MakeOrThrow<OpKernelInfoWrapper>(
+            ComPtr<OpKernelInfoWrapper> kernelInfoWrapper = Dml::SafeMakeOrThrow<OpKernelInfoWrapper>(
                 &Info(),
                 m_abiExecutionObject.Get(),
                 &inputShapes,
@@ -2569,7 +2569,7 @@ namespace Windows::AI::MachineLearning::Adapter
                 EdgeShapes localInferredOutputShapes;
                 ComPtr<IMLOperatorKernel> localKernel = inferShapesAndCreateKernel(local_input_shapes, localInferredOutputShapes);
 
-                ComPtr<OpKernelContextWrapper> kernelContextWrapper = wil::MakeOrThrow<OpKernelContextWrapper>(
+                ComPtr<OpKernelContextWrapper> kernelContextWrapper = Dml::SafeMakeOrThrow<OpKernelContextWrapper>(
                     context,
                     Info().GetExecutionProvider(),
                     m_internalOperator,
@@ -2588,7 +2588,7 @@ namespace Windows::AI::MachineLearning::Adapter
             }
         }
 
-        ComPtr<OpKernelContextWrapper> kernelContextWrapper = wil::MakeOrThrow<OpKernelContextWrapper>(
+        ComPtr<OpKernelContextWrapper> kernelContextWrapper = Dml::SafeMakeOrThrow<OpKernelContextWrapper>(
             context,
             Info().GetExecutionProvider(),
             m_internalOperator,
@@ -2811,7 +2811,7 @@ namespace Windows::AI::MachineLearning::Adapter
         onnxruntime::ProtoHelperNodeContext protoContext(node);
         onnxruntime::OpNodeProtoHelper<onnxruntime::ProtoHelperNodeContext> info(&protoContext);
 
-        ComPtr<MLKernelInferenceContext> inferenceContext = wil::MakeOrThrow<MLKernelInferenceContext>(&info, inputShapes, outputShapes, defaultAttributes, requiredConstantCpuInputs, constantInputGetter);
+        ComPtr<MLKernelInferenceContext> inferenceContext = Dml::SafeMakeOrThrow<MLKernelInferenceContext>(&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<IMLOperatorTensor> tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(
+                Microsoft::WRL::ComPtr<IMLOperatorTensor> tensorWrapper = Dml::SafeMakeOrThrow<OnnxTensorWrapper>(
                     const_cast<onnx::TensorProto*>(ctx->getInputData(index)), std::filesystem::path());
                 return tensorWrapper;
             }
         );
 
-        return wil::MakeOrThrow<MLSchemaInferenceContext>(info, ctx, requiredConstantCpuInputs, mlOperatorTensorGetter);
+        return Dml::SafeMakeOrThrow<MLSchemaInferenceContext>(info, ctx, requiredConstantCpuInputs, mlOperatorTensorGetter);
     }
 
     MLSchemaInferenceContext::MLSchemaInferenceContext(
@@ -2952,7 +2952,7 @@ namespace Windows::AI::MachineLearning::Adapter
         const AttributeMap* defaultAttributes)
     {
         MLOperatorTensorGetter mLOperatorTensorGetter = MLOperatorTensorGetter();
-        return wil::MakeOrThrow<MLSupportQueryContext>(info, defaultAttributes, mLOperatorTensorGetter);
+        return Dml::SafeMakeOrThrow<MLSupportQueryContext>(info, defaultAttributes, mLOperatorTensorGetter);
     }
 
     MLSupportQueryContext::MLSupportQueryContext(

onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlDFT.h

@@ -1097,7 +1097,7 @@ class GpuDFTOperatorFactory : public WRL::Base<IMLOperatorKernelFactory>
                 version = 20;
             }
 
-            auto dftOperator = wil::MakeOrThrow<GpuDFTOperator>(context, version);
+            auto dftOperator = Dml::SafeMakeOrThrow<GpuDFTOperator>(context, version);
             dftOperator.CopyTo(kernel);
             return S_OK;
         }
@@ -1177,8 +1177,8 @@ class GpuDFTOperatorFactory : public WRL::Base<IMLOperatorKernelFactory>
         kernelDescription.options = MLOperatorKernelOptions::None;
         kernelDescription.executionOptions = 0;
 
-        auto shareInferrer = wil::MakeOrThrow<DFTShapeInferrer>();
-        auto factory = wil::MakeOrThrow<GpuDFTOperatorFactory>();
+        auto shareInferrer = Dml::SafeMakeOrThrow<DFTShapeInferrer>();
+        auto factory = Dml::SafeMakeOrThrow<GpuDFTOperatorFactory>();
 
         std::array<uint32_t, 2> requiredConstantCpuInputs = { 1, 2 };
 

onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlGridSample.h

@@ -747,7 +747,7 @@ class DmlGridSampleOperatorFactory : public WRL::Base<IMLOperatorKernelFactory>
     {
         try
         {
-            auto dftOperator = wil::MakeOrThrow<DmlGridSampleOperator>(context);
+            auto dftOperator = Dml::SafeMakeOrThrow<DmlGridSampleOperator>(context);
             dftOperator.CopyTo(kernel);
             return S_OK;
         }
@@ -832,8 +832,8 @@ class DmlGridSampleOperatorFactory : public WRL::Base<IMLOperatorKernelFactory>
         kernelDescription.options = MLOperatorKernelOptions::None;
         kernelDescription.executionOptions = 0;
 
-        auto shareInferrer = wil::MakeOrThrow<GridSampleShapeInferrer>();
-        auto factory = wil::MakeOrThrow<DmlGridSampleOperatorFactory>();
+        auto shareInferrer = Dml::SafeMakeOrThrow<GridSampleShapeInferrer>();
+        auto factory = Dml::SafeMakeOrThrow<DmlGridSampleOperatorFactory>();
 
         ComPtr<IMLOperatorRegistryPrivate> registryPrivate;
         ORT_THROW_IF_FAILED(registry->QueryInterface(IID_PPV_ARGS(&registryPrivate)));

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<const uint32_t> outputShape
+        )
+    {
+        const uint32_t outputShapeDimCount = gsl::narrow_cast<uint32_t>(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<int32_t>(kernelInfo.GetAttribute<int64_t>(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<uint32_t>(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<uint32_t> 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<uint32_t>(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

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<const uint32_t> outputShape
+            );
+
         static void TryConvertTensorToBroadcastScalar(
             const MLOperatorKernelCreationContext& kernelInfo,
             const DML_TENSOR_DESC* tensor,