diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp index 9d0ba9dc7e..17aa197396 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlGraphFusionHelper.cpp @@ -292,7 +292,8 @@ namespace DmlGraphFusionHelper { for (size_t i = 0; i < graphDesc.nodes.size(); ++i) { - dmlOperatorGraphNodes[i] = DML_OPERATOR_GRAPH_NODE_DESC{graphDesc.nodes[i].op.Get()}; + auto& nodeInfo = graphDesc.nodes[i]; + dmlOperatorGraphNodes[i] = DML_OPERATOR_GRAPH_NODE_DESC{nodeInfo.op.Get(), nodeInfo.name.data()}; dmlGraphNodes[i] = DML_GRAPH_NODE_DESC{DML_GRAPH_NODE_TYPE_OPERATOR, &dmlOperatorGraphNodes[i]}; } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp index 6b9230657b..636f46428c 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp @@ -361,6 +361,7 @@ namespace Dml::GraphDescBuilder NodeInfo nodeInfo = {}; nodeInfo.op = std::move(op); + nodeInfo.name = node.Name(); graphNodes.push_back(std::move(nodeInfo)); } } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h index bf2ccae55e..5c04962e55 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h @@ -28,6 +28,7 @@ namespace Dml struct NodeInfo { Microsoft::WRL::ComPtr op; + std::string name; }; struct GraphDesc diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp index 197c62283f..af9755ac22 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp @@ -982,6 +982,88 @@ namespace Windows::AI::MachineLearning::Adapter m_abiExecutionObject.CopyTo(executionInterface); } + uint32_t STDMETHODCALLTYPE OpKernelInfoWrapper::GetUtf8NameBufferSizeInBytes() const noexcept + { + // Include null terminator. + return static_cast(m_impl->node().Name().size() + 1); + } + + HRESULT STDMETHODCALLTYPE OpKernelInfoWrapper::GetUtf8Name(uint32_t bufferSizeInBytes, char* outputName) const noexcept + { + if (bufferSizeInBytes == 0) + { + return E_INVALIDARG; + } + + // Copy as many characters as possible, leaving room for the null terminator. + const auto& nodeName = m_impl->node().Name(); + size_t charsCopied = nodeName.copy(outputName, bufferSizeInBytes - 1); + + // Write the null terminator. + assert(charsCopied >= 0 && charsCopied < bufferSizeInBytes); + outputName[charsCopied] = '\0'; + + return S_OK; + } + + uint32_t STDMETHODCALLTYPE OpKernelInfoWrapper::GetWideNameBufferSizeInBytes() const noexcept + { + const auto& name = m_impl->node().Name(); + if (name.empty()) + { + // Include null terminator. + return sizeof(wchar_t); + } + + int requiredSizeInChars = MultiByteToWideChar(CP_UTF8, 0, name.data(), static_cast(name.size()), nullptr, 0); + assert(requiredSizeInChars > 0); + + // Include null terminator. + return static_cast((requiredSizeInChars + 1) * sizeof(wchar_t)); + } + + HRESULT STDMETHODCALLTYPE OpKernelInfoWrapper::GetWideName(uint32_t bufferSizeInBytes, wchar_t* outputName) const noexcept + { + // Buffer needs to be large enough to at least hold a null terminator. + if (bufferSizeInBytes < sizeof(wchar_t)) + { + return E_INVALIDARG; + } + + const auto& nodeName = m_impl->node().Name(); + if (nodeName.empty()) + { + outputName[0] = L'\0'; + return S_OK; + } + + uint32_t bufferSizeInChars = bufferSizeInBytes / sizeof(wchar_t); + int charsCopiedIfSucceeded = MultiByteToWideChar(CP_UTF8, 0, nodeName.data(), static_cast(nodeName.size()), outputName, bufferSizeInChars); + + if (charsCopiedIfSucceeded > 0) + { + // The return value is only > 0 if ALL characters copied successfully. + // Write null terminator at the end of copied chars, which may not be at the end of the buffer. + outputName[charsCopiedIfSucceeded] = L'\0'; + return S_OK; + } + + // An error must have occurred in MultiByteToWideChar. + assert(charsCopiedIfSucceeded <= 0); + auto lastError = GetLastError(); + + if (lastError == ERROR_INSUFFICIENT_BUFFER) + { + // The buffer was too small, but MultiByteToWideChar will have copied as many chars as possible. + // Truncate and overwrite last char with null terminator. Don't treat this as an error. + outputName[bufferSizeInChars - 1] = L'\0'; + return S_OK; + } + + assert(lastError == ERROR_INVALID_PARAMETER || lastError == ERROR_NO_UNICODE_TRANSLATION); + return E_INVALIDARG; + } + template uint32_t STDMETHODCALLTYPE OpNodeInfoWrapper::GetInputCount() const noexcept { diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h index dd1b743587..ffb8bf1a38 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h @@ -331,7 +331,7 @@ class OnnxTensorWrapper : public WRL::Base, public Closable class OpKernelInfoWrapper : public OpNodeInfoWrapper< onnxruntime::ProtoHelperNodeContext, WRL::Base< - Microsoft::WRL::ChainInterfaces, + Microsoft::WRL::ChainInterfaces, IMLOperatorTensorShapeDescription, IMLOperatorAttributes1>, onnxruntime::null_type> { @@ -374,6 +374,14 @@ class OpKernelInfoWrapper : public OpNodeInfoWrapper< return E_NOTIMPL; } + // IMLOperatorKernelCreationContextNodeWrapperPrivate methods. + + uint32_t STDMETHODCALLTYPE GetUtf8NameBufferSizeInBytes() const noexcept override; + HRESULT STDMETHODCALLTYPE GetUtf8Name(uint32_t bufferSizeInBytes, char* name) const noexcept override; + + uint32_t STDMETHODCALLTYPE GetWideNameBufferSizeInBytes() const noexcept override; + HRESULT STDMETHODCALLTYPE GetWideName(uint32_t bufferSizeInBytes, wchar_t* name) const noexcept override; + private: // For shape info, in addition to the info const EdgeShapes* m_inferredOutputShapes = nullptr; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp index 3ae29629ef..497207ab83 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp @@ -89,6 +89,11 @@ namespace Dml { DML_EXECUTION_FLAGS executionFlags = GetExecutionFlags(); ORT_THROW_IF_FAILED(m_dmlDevice->CompileOperator(dmlOperator.Get(), executionFlags, IID_PPV_ARGS(&m_compiledOperator))); + + // Static buffer (might truncate name) to avoid excessive dynamic allocation only for debugging purposes. + wchar_t nodeName[512]; + ORT_THROW_IF_FAILED(kernelInfo.GetNodeWrapperInterface()->GetWideName(sizeof(nodeName), nodeName)); + ORT_THROW_IF_FAILED(m_compiledOperator->SetName(nodeName)); UINT64 persistentResourceSize = m_compiledOperator->GetBindingProperties().PersistentResourceSize; if (persistentResourceSize > 0) diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h index d79b2fb4e7..2e5de08f87 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h @@ -486,6 +486,7 @@ public: MLOperatorKernelCreationContext(IMLOperatorKernelCreationContext* impl) : MLOperatorAttributes(impl), m_impl(impl) { m_impl.As(&m_implPrivate); + m_impl.As(&m_nodeWrapperImpl); } // For cases of interop where the caller needs to pass the unwrapped class across a boundary. @@ -494,6 +495,11 @@ public: return m_impl; } + IMLOperatorKernelCreationContextNodeWrapperPrivate* GetNodeWrapperInterface() const noexcept + { + return m_nodeWrapperImpl.Get(); + } + Microsoft::WRL::ComPtr GetExecutionInterface() const noexcept { Microsoft::WRL::ComPtr ret; @@ -557,6 +563,7 @@ public: private: Microsoft::WRL::ComPtr m_impl; Microsoft::WRL::ComPtr m_implPrivate; + Microsoft::WRL::ComPtr m_nodeWrapperImpl; }; class MLShapeInferenceContext : public MLOperatorAttributes diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorPrivate.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorPrivate.h index afcfc587ec..6d906c6056 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorPrivate.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorPrivate.h @@ -57,6 +57,30 @@ IMLOperatorKernelCreationContextPrivate : public IMLOperatorKernelCreationContex ) const noexcept PURE; }; +interface __declspec(uuid("1d2e1226-a918-4236-8775-175cf1f52c9a")) +IMLOperatorKernelCreationContextNodeWrapperPrivate : public IMLOperatorKernelCreationContextPrivate +{ + //! Gets the minimum size of a char buffer to store the node name (including null terminator). + //! Returns 1 if the node has no name (calling GetUtf8Name will write a single null terminator). + STDMETHOD_(uint32_t, GetUtf8NameBufferSizeInBytes)() const noexcept PURE; + + //! Writes the node name and null terminator into a char buffer. + STDMETHOD(GetUtf8Name)( + uint32_t bufferSizeInBytes, + _Out_writes_(bufferSizeInBytes) char* name + ) const noexcept PURE; + + //! Gets the minimum size of a wchar buffer to store the node name (including null terminator). + //! Returns sizeof(wchar_t) if the node has no name (calling GetWideName will write a null terminator). + STDMETHOD_(uint32_t, GetWideNameBufferSizeInBytes)() const noexcept PURE; + + //! Writes the node name and null terminator into a wchar buffer. + STDMETHOD(GetWideName)( + uint32_t bufferSizeInBytes, + _Out_writes_(bufferSizeInBytes) wchar_t* name + ) const noexcept PURE; +}; + //! \interface IMLOperatorAttributes1 //! \brief Represents the values of an operator's attributes, as determined by a model using the operator. //! This interface is called by implementations of custom operator kernels, and by implementations