From 03cfb7d73ede48345d90688071232bd76e813e1f Mon Sep 17 00:00:00 2001 From: Justin Stoecker Date: Thu, 2 Feb 2023 13:42:15 -0800 Subject: [PATCH] Use ORT node names in DML graphs/ops (#14461) ### Description Applies ORT node names to corresponding compiled operators or DML graph nodes. ### Motivation and Context This makes it easier to correlate ONNX nodes to events in PIX GPU captures when using the DML EP. Names set in the DML graph nodes require additional modifications to the DML runtime library (available in a future NuGet package). --- .../src/DmlGraphFusionHelper.cpp | 3 +- .../src/GraphDescBuilder.cpp | 1 + .../src/GraphDescBuilder.h | 1 + .../src/MLOperatorAuthorImpl.cpp | 82 +++++++++++++++++++ .../src/MLOperatorAuthorImpl.h | 10 ++- .../src/Operators/DmlOperator.cpp | 5 ++ .../MLOperatorAuthorHelper.h | 7 ++ .../MLOperatorAuthorPrivate.h | 24 ++++++ 8 files changed, 131 insertions(+), 2 deletions(-) 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