mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-20 19:12:24 +00:00
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:
parent
0bcca7ad45
commit
03cfb7d73e
8 changed files with 131 additions and 2 deletions
|
|
@ -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]};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -361,6 +361,7 @@ namespace Dml::GraphDescBuilder
|
|||
|
||||
NodeInfo nodeInfo = {};
|
||||
nodeInfo.op = std::move(op);
|
||||
nodeInfo.name = node.Name();
|
||||
graphNodes.push_back(std::move(nodeInfo));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ namespace Dml
|
|||
struct NodeInfo
|
||||
{
|
||||
Microsoft::WRL::ComPtr<IDMLOperator> op;
|
||||
std::string name;
|
||||
};
|
||||
|
||||
struct GraphDesc
|
||||
|
|
|
|||
|
|
@ -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
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue