saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 04:39:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix DML custom operators which set descriptor heap to command list (#12059)

Browse source
This commit is contained in:
Jeff Bloomfield 2022-07-01 09:49:23 -07:00 committed by GitHub
parent bfe1eca10c
commit 02b9b12127
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 18 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/core/providers/dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h
View file

@ -45,7 +45,7 @@ namespace Windows::AI::MachineLearning::Adapter
IUnknown* data,
bool isInternalOperator) = 0;
virtual void GetABIExecutionInterface(
virtual void GetABIExecutionInterfaceAndInvalidateState(
bool isInternalOperator,
IUnknown** abiExecutionObject) const = 0;

6
onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.h
View file

@ -61,6 +61,12 @@ namespace Dml
return m_operationsRecordedInCurrentCommandList || !m_pendingCommandLists.empty();
}
// Forces the descriptor heap to be reset to D3D before executing future operations
void InvalidateDescriptorHeap()
{
m_currentDescriptorHeap = nullptr;
}
private:
std::shared_ptr<CommandQueue> m_queue;

6
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.cpp
View file

@ -125,10 +125,14 @@ namespace Dml
m_dmlRecorder.ResourceBarrier(barriers);
}
void ExecutionContext::GetCommandListForRecording(ID3D12GraphicsCommandList** commandList)
void ExecutionContext::GetCommandListForRecordingAndInvalidateState(ID3D12GraphicsCommandList** commandList)
{
assert(!m_closed);
SetCommandRecorder(&m_dmlRecorder);
// Ensure the descriptor heap is reset to D3D as something external may change it before recording
m_dmlRecorder.InvalidateDescriptorHeap();
m_dmlRecorder.GetCommandList().CopyTo(commandList);
}

2
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.h
View file

@ -66,7 +66,7 @@ namespace Dml
void AddUAVBarrier();
void ResourceBarrier(gsl::span<const D3D12_RESOURCE_BARRIER> barriers);
void GetCommandListForRecording(ID3D12GraphicsCommandList** commandList);
void GetCommandListForRecordingAndInvalidateState(ID3D12GraphicsCommandList** commandList);
// Forces all queued work to begin executing on the GPU. This method returns immediately and does not wait
// for the submitted work to complete execution on the GPU.

4
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
View file

@ -688,7 +688,7 @@ namespace Dml
return m_allocator->DecodeDataHandle(data)->GetPooledResourceId();
}
void ExecutionProviderImpl::GetABIExecutionInterface(
void ExecutionProviderImpl::GetABIExecutionInterfaceAndInvalidateState(
bool isInternalOperator,
IUnknown** abiExecutionObject) const
{
@ -702,7 +702,7 @@ namespace Dml
else
{
ComPtr<ID3D12GraphicsCommandList> commandList;
m_context->GetCommandListForRecording(commandList.GetAddressOf());
m_context->GetCommandListForRecordingAndInvalidateState(commandList.GetAddressOf());
#ifdef _GAMING_XBOX
ComPtr<GraphicsUnknownWrapper> wrappedCommandList = Microsoft::WRL::Make<GraphicsUnknownWrapper>(commandList.Get());
*abiExecutionObject = wrappedCommandList.Detach();

2
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h
View file

@ -108,7 +108,7 @@ namespace Dml
IUnknown* data,
bool isInternalOperator) override;
void GetABIExecutionInterface(
void GetABIExecutionInterfaceAndInvalidateState(
bool isInternalOperator,
IUnknown** abiExecutionObject) const override;

4
onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp
View file

@ -1467,7 +1467,7 @@ namespace Windows::AI::MachineLearning::Adapter
m_abiExecutionObject = m_providerExecutionObject;
if (m_winmlProvider)
{
m_winmlProvider->GetABIExecutionInterface(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf());
m_winmlProvider->GetABIExecutionInterfaceAndInvalidateState(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf());
}
TransitionResourcesForOperatorIfRequired(true);
@ -1737,7 +1737,7 @@ namespace Windows::AI::MachineLearning::Adapter
if (m_winmlProvider)
{
// Get the particular object to return to a isInternalOperator based on the registration of that kernel.
m_winmlProvider->GetABIExecutionInterface(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf());
m_winmlProvider->GetABIExecutionInterfaceAndInvalidateState(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf());
}
}