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).
This commit is contained in:
Justin Stoecker 2023-02-02 13:42:15 -08:00 committed by GitHub
parent 0bcca7ad45
commit 03cfb7d73e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 131 additions and 2 deletions

View file

@ -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]};
}

View file

@ -361,6 +361,7 @@ namespace Dml::GraphDescBuilder
NodeInfo nodeInfo = {};
nodeInfo.op = std::move(op);
nodeInfo.name = node.Name();
graphNodes.push_back(std::move(nodeInfo));
}
}

View file

@ -28,6 +28,7 @@ namespace Dml
struct NodeInfo
{
Microsoft::WRL::ComPtr<IDMLOperator> op;
std::string name;
};
struct GraphDesc

View file

@ -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<uint32_t>(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<int>(name.size()), nullptr, 0);
assert(requiredSizeInChars > 0);
// Include null terminator.
return static_cast<uint32_t>((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<int>(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 <class NodeInfoImpl_t, class Base1_t, class Base2_t>
uint32_t STDMETHODCALLTYPE OpNodeInfoWrapper<NodeInfoImpl_t, Base1_t, Base2_t>::GetInputCount() const noexcept
{

View file

@ -331,7 +331,7 @@ class OnnxTensorWrapper : public WRL::Base<IMLOperatorTensor>, public Closable
class OpKernelInfoWrapper : public OpNodeInfoWrapper<
onnxruntime::ProtoHelperNodeContext,
WRL::Base<
Microsoft::WRL::ChainInterfaces<IMLOperatorKernelCreationContextPrivate, IMLOperatorKernelCreationContext>,
Microsoft::WRL::ChainInterfaces<IMLOperatorKernelCreationContextNodeWrapperPrivate, IMLOperatorKernelCreationContextPrivate, IMLOperatorKernelCreationContext>,
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;

View file

@ -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)

View file

@ -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<IUnknown> GetExecutionInterface() const noexcept
{
Microsoft::WRL::ComPtr<IUnknown> ret;
@ -557,6 +563,7 @@ public:
private:
Microsoft::WRL::ComPtr<IMLOperatorKernelCreationContext> m_impl;
Microsoft::WRL::ComPtr<IMLOperatorKernelCreationContextPrivate> m_implPrivate;
Microsoft::WRL::ComPtr<IMLOperatorKernelCreationContextNodeWrapperPrivate> m_nodeWrapperImpl;
};
class MLShapeInferenceContext : public MLOperatorAttributes

View file

@ -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